Codeforces Round #225 (Div. 1) D. Antimatter

こんな簡単に解けると思ってなかったので目からウロコでした。

問題

codeforces.com

解法

dp[n][sum] = ([1, n], [2, n], ..., [n-1, n] という区間で, 和が sum になる場合の数) というのを考えます(antimatter の分はマイナスで数えるので, 要するに dp[n][0] の和が答えになります。ただ, コード上 dp[5][-4] とかは出来ないので下のコードでは MAX 分だけオフセットをとっていると)。

まず, 普通に dp[n-1][sum] を使って, dp[n][sum+a], dp[n][sum-a] に場合の数を伝搬させる普通の dp をやります。で, あとは n 番目からのスタートする区間のことを考えないといけないですが, これは dp[n][-a]++, dp[n][a]++ とやれば「n 番目から初めて合計 a になるのは 1 通り」というのを考慮しているのでこれで OK です。頭良いですね。

const ll MOD = 1e9+7;
const int MAX = 10000;
ll dp[2][2*MAX+1];

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    int n, cur = 0, next = 1;
    ll ans = 0;
    cin >> n;
    for (int i = 0; i < n; i++) {
        int a;
        cin >> a;
        memset(dp[next], 0, sizeof(dp[next]));
        for (int i = 0; i <= 2*MAX; i++) if (dp[cur][i]) {
            (dp[next][i+a] += dp[cur][i]) %= MOD;
            (dp[next][i-a] += dp[cur][i]) %= MOD;
        }
        ans += dp[next][MAX];
        dp[next][MAX+a]++;
        dp[next][MAX-a]++;
        swap(cur, next);
    }
    cout << ans%MOD << endl;
    return 0;
}