AtCoder Educational DP Contest T問題解いてみた

今回の問題

問題文

$N$ を正整数とします。 長さ $N-1$ の文字列 $s$ が与えられます。
$s$ は < と > からなります。

$(1, 2, \ldots, N)$ を並べ替えた順列 $(p_1, p_2, \ldots, p_N)$ であって、次の条件を満たすものは何通りでしょうか？
$10^9 + 7$ で割った余りを求めてください。

• 各 $i$ ($1 \leq i \leq N-1$) について、$s$ の $i$ 文字目が < の場合は $p_i < p_{i+1}$ であり、$s$ の $i$ 文字目が `>` の場合は $p_i > p_{i+1}$ である。

制約

• $N$ は整数である。
• $2 \leq N \leq 3000$
• $s$ は長さ $N-1$ の文字列である。
• $s$ は < と > からなる。

回答

#include <bits/stdc++.h>
using namespace std;

const int MOD = 1e9 + 7;

// dp[i][j]
// i: 現在の順列の長さ (1 ~ N)
// j: 末尾の数が、現在使っている数の中で j 番目に小さい (1 ~ i)
// N <= 3000 なので 2次元配列でOK
int dp[3005][3005];

int main() {
    int N;
    string s;
    cin >> N >> s;

    // 初期化：長さ1の順列は「1」のみ。1番目に小さい数は1つだけ。
    dp[1][1] = 1;

    // 長さ i = 2 から N まで順番に計算
    for (int i = 2; i <= N; i++) {
        
        // 高速化のため、一つ前の段 (dp[i-1]) の累積和を作る
        vector<int> S(i + 1, 0); // S[k] = dp[i-1][1] + ... + dp[i-1][k]
        for (int k = 1; k <= i - 1; k++) {
            S[k] = (S[k - 1] + dp[i - 1][k]) % MOD;
        }

        // 今回の文字 s[i-2] を見て条件分岐 (0-indexedなのでズレに注意)
        // 文字列 s の k 文字目は、p_{k+1} と p_{k+2} の関係を表す
        // s[i-2] は p_{i-1} と p_i の関係
        
        for (int j = 1; j <= i; j++) {
            if (s[i - 2] == '<') {
                // p_{i-1} < p_i なので、前の順位 k は j より小さい (k < j)
                // dp[i][j] = sum(dp[i-1][1] ... dp[i-1][j-1])
                // 累積和 S を使えば S[j-1]
                dp[i][j] = S[j - 1];
            } else {
                // p_{i-1} > p_i なので、前の順位 k は j 以上 (k >= j)
                // dp[i][j] = sum(dp[i-1][j] ... dp[i-1][i-1])
                // 累積和 S を使えば S[i-1] - S[j-1]
                int val = (S[i - 1] - S[j - 1] + MOD) % MOD; // 引き算は負の数対策
                dp[i][j] = val;
            }
        }
    }

    // 答えは dp[N][1] ... dp[N][N] の総和
    long long ans = 0;
    for (int j = 1; j <= N; j++) {
        ans = (ans + dp[N][j]) % MOD;
    }

    cout << ans << endl;

    return 0;
}

自分なりの解説

重要なのは、具体的な数の大きさではなく相対的な数の大きさであることに気づけるかが重要だった。
dp[i][j]を長さ$i$の順列で、末尾の項$p_i$が、そこまでに使った $i$ 個の数の中で $j$ 番目に小さい数 であるような並べ方の総数とする。
また、このままでは間に合わないので累積和を使ってさらに高速化している。