Colorful Hats 2 [Sumitomo Mitsui Trust Bank Programming Contest 2019 E]

https://atcoder.jp/contests/sumitrust2019/tasks/sumitb2019_e

解説

https://atcoder.jp/contests/sumitrust2019/submissions/8777300

着色をしていくわけだが、先頭から順番に色をつけていこう。
最初の0は3色のどれかが入る。次の0は残りの2色。最後は1色しか無い。
先頭から1が出てきたら、それより前の0の個数分だけの選択肢がある。
次に1が出てきたら、0は1つ使われているので、それより前の0の個数-1分だけの選択肢がある。
この選択肢は全部独立に考えることができるので、組み合わせをかけ合わせる。
足りなくなったりしたら、答えは0。
先頭から順番に個数を数えて、減らしながらやっていくと答え。
実装を楽にするためにAの全要素+1しておき、0が3つ既にあるものと数えると、最初をあまり考えなくていい。

int N, A[101010];
int cnt[101010];
//---------------------------------------------------------------------------------------------------
mint solve() {
    mint ans = 1;
    rep(i, 0, N) A[i]++;
    cnt[0] = 3;
    rep(i, 0, N) {
        if (0 < cnt[A[i] - 1]) {
            ans *= cnt[A[i] - 1];
            cnt[A[i] - 1]--;
            cnt[A[i]]++;
        }
        else return 0;
    }
    return ans;
}
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N;
    rep(i, 0, N) cin >> A[i];
    cout << solve() << endl;
}