三井住友信託銀行プログラミングコンテスト2019 E – Colorful Hats 2 をPython3で解く

Share

AtCoder上にある問題のうち、AtCoder Problemsでdiff 800以上と判定されているものを順番に解いていく企画。
基本的な考え方は全てコード中のコメントに入れてあるので、参照のこと。

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

二次元配列をどんどん更新していくDP。遷移自体はそこまで複雑でもないので、落ち着いて考えれば解けると思う。

# 三井住友信託銀行プログラミングコンテスト2019 E - Colorful Hats 2
# https://atcoder.jp/contests/sumitrust2019/tasks/sumitb2019_e
# tag: 考察 MOD 順列・組み合わせ 数え上げ

# 入力例 3 を眺めると気づきやすいが、出来上がる数列は
# RGB に対応した 3 つの要素に分解できる。

# 考えてみれば当たり前の話で、数列の中から同じ色の帽子を
# かぶっている人だけを抽出すれば、その数列は
# [0, 1, 2, 3...] と連番になっていなければならないため。

# すなわち、最初は [0, 0, 0] という 3 つの要素があるとして、
# このうちどれかを数列に加え、選んだ要素に 1 を足すという
# 操作を繰り返すような形になっている。

# この全体での通り数については、各要素を選ぶ際に、
# 数列に対応した帽子がその時点で何種類あるかを考えることで
# 計算できる。

# 具体的には、

# [0, 0, 0]
# 最初の数字が 0 なら、3 つのうちから一つを選ぶ。

# → [0, 0, 1]
# 次の数字が 1 なら、選べる帽子は一つしか無い。

# → [0, 0, 2]
# 次の数字が 0 なら、2 つのうちから一つを選ぶ。

# → [0, 1, 2] ...

def main():
    N = int(input())
    A = list(map(int, input().split()))
    MOD = 10**9+7

    nums = [0, 0, 0]

    result = 1

    for a in A:
        # 要素は 3 個しか無いので、線形探索で問題ない。
        mult = nums.count(a)
        result = (result * mult) % MOD

        nums[nums.index(a)] += 1

    print(result)

main()
Share

コメントを残す

メールアドレスが公開されることはありません。