AtCoder-ABC210 C - Colorful Candies【Python解答例】

AtCoder Beginner Contest210のC問題についてPythonの解答例を記事にしていきます。
AtCoder Beginner Contest 210 - AtCoder

AtCoder Beginner Contest210 C - Colorful Candies

C - Colorful Candies

問題文

N 個のキャンディが左右一列に並んでいます。
それぞれのキャンディは、色 1、色 2、…、色 10^9の、10^9 種類の色のうちいずれかの色をしています。
i=1,2,…,N について、左から i 番目のキャンディの色は色 ci です。

高橋君は並んでいるキャンディのうち、連続して並んだ K 個のキャンディをもらうことができます。
すなわち、1≤i≤N−K+1 を満たす整数 i を選んで、 左から i 番目、i+1 番目、i+2 番目、…、i+K−1 番目のキャンディをもらうことができます。

高橋君はいろいろな色のキャンディを食べたいので、 もらうキャンディに含まれる色の種類数が多いほどうれしい気持ちになります。
高橋君がもらうキャンディに含まれる色の種類数の最大値を出力してください。

制約

・1≤K≤N≤3×10^5
・1≤ci≤10^9
・入力はすべて整数

解答例

n, k = map(int,input().split())

C = list(map(int, input().split()))

color = {}

for i in range(n):
    color.setdefault(C[i], 0)

tmp = 0
for i in range(k):
    if color[C[i]] == 0:
        tmp += 1
    color[C[i]] += 1

ans = tmp
for i in range(n-k):
    if color[C[i]] == 1:
        tmp -= 1
    color[C[i]] -= 1
    if color[C[i+k]] == 0:
        tmp += 1
    color[C[i+k]] += 1
    ans = max(ans, tmp)

print(ans)

解説

左からi番目のキャンディが色Ciとして、連続したK個のキャンディをもらうときに最大何色のキャンディをもらえるか答える問題です。

キャンディの色を辞書colorで管理していきます。

まず、キャンディの色をすべて辞書にキーとして登録しました。
tmpをある区間においてK個のアメをもらときの色の数とします。
ある色においてcolorの要素の数を調べて、要素数が0の新しい色が加わる場合はtmpに1を足します。
要素数が1の色が区間から抜ける場合はtmpから1を引きます。
このようにtmpを管理することですべての区間における色の数を把握することができました。

最後にtmpの最大値ansを解答することでACでした。

ABC210の関連記事はこちら
ebisuke33.hatenablog.com
ebisuke33.hatenablog.com