CS Academy Restricted Permutations

リンク : https://csacademy.com/contest/archive/task/restricted-permutations/

問題概要

Nが与えられ、N-1個の数字が与えられる。この数字はi番目の数字が1ならば、iより後ろにi+1が存在する、0ならばiより前にi+1が存在する。N-1個の数字が示す条件をすべて満たす長さNの順列の通り数を求めよ。

解説

例えば、(1,1,1,0)なら下の図のように考えられる。どこに4を入れるかのみを考えればよい。 この図の通り考えると、通り数をも求めるだけなら、前から何番目に位置するかが重要であるように見える。そこで、dp[i][j]をiまでを使い、iまでの条件を満たし、iが前からj番目にある順列の通り数とすると、
i番目の条件が1の場合は
dp[i][j+1]=sum{k=0,1...j}{dp[i-1][k]}
i番目の条件が2の場合は
dp[i][j]=sum{k=j,j+1...i}{dp[i-1][k]}
となる。sumの中の部分は累積和をもっておくとO(1)で計算できるためO(n²)でこの問題が解ける。
余談だがO(n² logn)を最初に提出し、TLEした...

LL n, ans,a[2100],dp[2100][2100],sum[2100];
int main() {
    cin >> n;
    rep(i, n - 1) {
        cin >> a[i];
    }
    dp[0][0] = 1;
    rep(i, n - 1) {
        rep(j, n) {
            sum[j + 1] = dp[i][j] + sum[j];
            sum[j + 1] %= MOD;
        }
        for (int j = 0; j <= i; j++) {
            if (a[i]) {
                dp[i + 1][j + 1] = sum[j + 1];
                dp[i + 1][j + 1] %= MOD;
            }
            else {
                dp[i + 1][j] = sum[n] - sum[j];
                dp[i + 1][j] = (dp[i + 1][j] + MOD) % MOD;
            }
        }
    }
    rep(i, n) {
        ans += dp[n-1][i];
        ans %= MOD;
    }
    cout << ans << endl;
    return 0;
}