JOI2011予選4 A First Grader （競プロ#1）

 最近競技プログラミングに取り組み始めたので練習日記のような形で解説記事を投稿していこうと思います。未熟者ですのでご指摘等あれば遠慮なくお申し付けください。

ここ数日は動的計画法の勉強を集中的に行っているのでまずはそのうちの一問をとりあげます。

方針

愚直に+と-の組み合わせを列挙するのはO( 2^N)なので無理だということはすぐにわかる。そこで動的計画法による解法を考える。動的計画法について学んだことがないという方は以下のdrkenさんの記事がとてもわかりやすいので読んでみることをおすすめする。

動的計画法超入門！ Educational DP Contest の A ～ E 問題の解説と類題集 - Qiita

今回の問題の場合はどうすればよいか、サンプル1を見ながら考えてみよう。

11
8 3 2 4 8 7 2 4 0 8 8

これは最初の10個の数字を足し算引き算した結果として8になるようなものが何個あるか、という問題だ。この時最初の10個の計算結果が8になるには最初の9個の計算結果が16か0になっていなければならないことが分かる。当然16-8=8, 0+8=8だからである。すなわち10個から8を与える組み合わせは9個から16を与える組み合わせと9個から0を与える組み合わせの和であることが分かる。ここまでくれば動的計画法が有効だということはすぐに分かるだろう。

よって最初からi個の数字を使った計算結果がjになる組み合わせを dp[i][j] として、i+1番目まで使った結果dp[i+1][j] は次のような式で表せる。ただしここではi+1番目の数が a[i+1] と直感的にわかりやすいように入力配列aを1インデックスとしているが実装上は a[i] としている。

dp[i+1] [j ] = dp[i][j+a[i+1]] + dp[i][j-a[i+1]]

この漸化式に基づいて更新していけばよい。ただし気をつけることとして制約上途中で0未満あるいは20より大きい数になってはいけないので更新の際に実装例のように条件を課してやる必要がある。

また自分は最初安易にdp配列の初期化をdp[0][0] = 1としていたが、これは入力の最初が0だった場合に出力が2倍になってしまうので注意していただきたい。

コード

#include <iostream>
#include <vector>
using namespace std;

int main() {
    int N; cin >> N;
    vector<int> a(N);
    for (int i = 0; i < N; ++i) cin >> a[i];
    long long dp[110][25] = {0};
    dp[1][a[0]] = 1;
    for (int i = 1; i < N-1; ++i) {
        for (int j = 0; j <= 20; ++j) {
            if (j + a[i] <= 20) dp[i+1][j] += dp[i][j+a[i]];
            if (j - a[i] >= 0) dp[i+1][j] += dp[i][j-a[i]];
        }
    }
    cout << dp[N-1][a[N-1]] << endl;
}