みんなのプロコン 2019 E問題：Pass

問題. Pass

$n$ 人のすぬけ君がそれぞれボールを2個ずつ持って1列に並んでいる．ボールの色はそれぞれ赤か青のどちらかである．各ステップで自分の持っているボールをどれか1つ前の人に渡すというステップを $2 n$ 回行う．先頭の人が選んだボールの色の列としてありうるものの個数を求めよ．

制約： $2 \le n \le 2,000$

解法．動的計画法

$r_i$ と $b_i$ をそれぞれ先頭から $i$ 番目までの人達が持っている赤色と青色のボールの総数とする．ただし便宜上，任意の $j \in (n, 2n]$ に対して $r_j$ と $b_j$ の値はそれぞれ $r_n$ と $b_n$ と等しいとする．
先頭の人が選んだボールの色の列を $s_1, s_2, \ldots, s_{2 n}$ としたとき，各 $i \in [1, 2n]$ に対して $s_1, s_2, \ldots, s_i$ で選んだ赤色と青色のボールの総数はそれぞれ $r_i, b_i$ 以下でないといけない．そのようなものの個数を動的計画法で求める．具体的には dp[i][j] を，先頭の人がボールを $i$ 個選んだときに，赤色のボールが $j$ 個含まれるようなものの数として漸化式を立てる．このとき， $j \le r_i$ を満たし，選ばれた青色のボールの個数が $i - j$ であることから $i - j \le b_i$ を満たすように遷移を考える（詳しくはソースコードを参照）．

計算時間： $O(n^2)$

ソースコードを表示

#include <bits/stdc++.h>

using namespace std;

using ll = long long;
constexpr ll MOD = 998244353;

ll Solve() {
    string s;
    cin >> s;

    const int n = s.size();
    vector<int> b(2 * n + 1), r(2 * n + 1);
    for (int i = 1; i <= n; ++i) {
        if (s[i - 1] == '0') r[i] = 2;
        else if (s[i - 1] == '2') b[i] = 2;
        else { ++r[i]; ++b[i]; }
    }
    for (int i = 1; i <= 2 * n; ++i) {
        r[i] += r[i - 1];
        b[i] += b[i - 1];
    }

    vector<vector<ll>> dp(2 * n + 1, vector<ll>(r[2 * n] + 1));
    dp[0][0] = 1;
    for (int i = 1; i <= 2 * n; ++i) {
        for (int j = max(0, i - b[i]); j <= r[i] && 0 <= i - j; ++j) {
            if (0 < j) (dp[i][j] += dp[i - 1][j - 1]) %= MOD;
            (dp[i][j] += dp[i - 1][j]) %= MOD;
        }
    }
    return dp[2 * n][r[2 * n]];
}

int main() {
    cin.tie(0); ios::sync_with_stdio(false);

    cout << Solve() << endl;

    return 0;
}

コンテスト中は問題を確認してなかったけど，コンテスト後に見たら解けたのでD問題よりもE問題を解くべきだったなと反省．