[AtCoder] 三井住友信託銀行プログラミングコンテスト2019 E – Colorful Hats 2

2020年10月19日2020年12月12日

問題

方針

\( A_1 \neq 0 \) のとき場合の数は \( 0 \) となります。ここで、\( d_i(j) \) を \( i \) 番目までの人を調べ終えた時、同じ帽子が \( j \) 個であるパターン数とします。初期値は、\( d_0(0) = 3 \) です。また、\( i \) 番目までの人を調べ終えた時の場合の数を \( p_i \) とします。初期値は \( p_0 = 1 \) です。\( A_1 = 0 \) であるとき、\( p_1 = p_0d_0(A_1) \) となり、\( d_1 \leftarrow d_0 \) として、

\begin{eqnarray}
d_1(A_1) & \leftarrow & d_1(A_1) – 1\\
d_1(A_1 + 1) & \leftarrow & d_1(A_1 + 1) + 1
\end{eqnarray}

と更新されます。\( d_i \) と \( d_{i + 1} \) で変わる値は \( 2 \) 箇所なので、\( d_i(j) \) を \( d(i) \) と再定義して使いまわすことにします。したがって、

\begin{eqnarray}
p_i &=& p_{i – 1}d(A_i) \\
d(A_i) & \leftarrow & d(A_i) – 1\\
d(A_i + 1) & \leftarrow & d(A_i + 1) + 1
\end{eqnarray}

となります。\( i \) 番目の人がある帽子を被ると、同じ帽子が \( A_i \) 個であるパターンが \( 1 \) 減り、 同じ帽子が \( A_i + 1\) 個であるパターンが \( 1 \) つ増えます。

コード

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

const ll mod = 1000000007;

int main() {
    int N;
    cin >> N;
    int A[N];
    for (int i = 0; i < N; i++) {
        cin >> A[i];
    }
    if (A[0] != 0) {
        cout << "0\n";
        return 0;
    }
    ll d[N + 1]{}; // d[i] 同じ帽子がi個であるときパターン数
    d[0] = 2;
    d[1] = 1;
    ll ans = 3;
    for (int i = 1; i < N; i++) {
        ans = ans * (d[A[i]]);
        ans %= mod;
        d[A[i]]--;
        d[A[i] + 1]++;
    }
    cout << ans << "\n";
    return 0;
}

参考

感想

こういう問題苦手です。解説読んでもあまり理解できませんでした。