三井住友信託銀行コン2019 E - Colorful Hats 2

問題原文

E - Colorful Hats 2

問題要旨

人の人が一列に並んでおり、それぞれ赤, 青, 緑 のいずれかの色の帽子を被っている。また各人は、「自分と同じ色の帽子を 被っている人が、前に 人いる」と主張している。これらの主張を全て満たすような帽子の色の組み合わせは何通りあるか。

考えたこと

• 数え上げと [tex: 10⁹ + 7] が並んでると小難しいDPを考えてしまう。。。

考えるべきこと

• 「今それぞれの色が何回出現したか？」を覚えておきたい
• 色は関係ないため、「 回出現した色が何色あるか？」を覚えておきたい
• ここで、 0回出現した色は3色ある と考える。
• すると、 ある が 仮に３だとすると、それは「3回出現した色のうちどれでも好きに選び」「3回出現した色の数を1減らし」「4回出現した色の数を1増やす」 という操作になる。

実装

from collections import defaultdict
N = int(input())
A = list(map(int, input().split()))
MOD = 10 ** 9 + 7
D = defaultdict(int)
# 0回出現した色は3色すべて
D[0] = 3

ans = 1
for a in A:
    # 選べるもののうち好きに色を選んで良い
    ans *= D[a]
    ans %= MOD
    
    # ある色がn回出現 → n+1回出現になる
    D[a] -= 1
    D[a + 1] += 1

print(ans % MOD)