競プロ初中級者100問

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

はじめに

カテゴリー競プロ初中級者100問では、Qiitaにて@e869120さんがレッドコーダーが教える、競プロ・AtCoder上達のガイドライン【中級編：目指せ水色コーダー！】としてまとめられている100問をPythonで解いています。

全問題の一覧はこちらです

問題

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

方針

Aを先頭から順番に色を割り当てていきます。この時、それぞれの色が割り当てられている人数を保持しきます。

複数の色に割り当てられる人がいますが、この場合は割り当てられる色のうち任意の色としてしまっても一般性を失いません。ただし、割り当て可能な人数を解（組み合わせ）にかけています。

コード

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
N = int(input())
A = list(map(int, input().split()))
MOD = 1000000007
cnt = [0,0,0] # 現在割り当てられている人数
ans = 1 # 現時点での組み合わせ
for a in A:
    idx = []
    for i in range(3):
        if cnt[i] == a: # 割り当て可能
            idx.append(i)
    if len(idx)==0: # 割り当て不能
        print(0)
        exit()
    cnt[idx[0]]+=1 # 任意の色で良いので、一番番号が若い色とする
    ans *= len(idx) # 割り当て可能な人数をかける
    ans %= MOD
print (ans)

記事情報

  • 投稿日：2020年6月4日
  • 最終更新日：2020年6月4日