三井住友信託銀行プログラミングコンテスト2019 E. Colorful Hats 2
問題
N人が一列に並んでおり、各人は赤、青、緑いずれかの帽子をかぶっている。
そして、左からi番目の人は「自分より左で、自分と同じ色の帽子をかぶっている人はAi人いる」と言っている情報が与えられる。
この情報に基づいた場合、帽子の色の組み合わせとして考えられるのは何通りか、mod1000000007で答えよ。
制約
- 1 <= N <= 100000
- 0 <= Ai <= N-1
解法
0 -> 1 -> 2 -> 3 -> ...とできる取り方は何通りか？という問題と考えられる。
dp的に考えようとすると、ある地点で取れる組み合わせは色を無視すると一意に決まることがわかる。
(「1まで出ているのが2個＆2まで出ているのが1個」のようなもの)
ある地点でAiだった場合、Ai-1までのものがすでに1つ以上出ているはずなので、これがx個だった場合は、Ai-1のx個のうちから1つ選んでAiにするようにすればよく、組み合わせはx倍になる。
これを初期値を1通りとして、上記を繰り返せばよい。
また、色の組み合わせとして、0が1回なら3通り、0が2回または3回ならば6通りの割当があり得るので、それをかければ良い。
この問題では、不正な入力もあり得る。
0が4個以上ある時や、最初が0以外で始まる場合は0通りを返す。
また、途中で不正な入力だった場合は、Ai-1の個数が0個になっているはずなので、0通りを返す。
