[백준 15961] 회전 초밥(Python)

1. 아이디어

투 포인터와 슬라이딩 윈도우를 이용해 해결한다. 처음 딱 1번, 연속된 k개 초밥을 선형탐색하여 초밥의 종류 수와 해당 종류를 몇개 먹었는지 카운트한다. 이후 lp, rp를 사용할 것이기 때문에, 반복문이 필요없다.

lp와 rp를 각각 1증가 시키기전에 lp에서 먹은 초밥 개수를 1 차감한다. 이 때, 먹은 초밥의 개수가 0이 되면, 초밥 종류 수를 1 차감한다. 이후 lp, rp를 각 1씩 증가시킨다. rp에있는 초밥 종류를 처음 먹을 때, 초밥 종류수를 1 증가시킨다. 그리고 이 종류에 대한 먹은 초밥 수도 1 증가시킨다.

한편, 구간[lp, rp]에서 쿠폰이 적용되는 초밥 종류가 없으면 초밥 종류 수를 1 증가시킨다.

위의 연산을 lp가 n이 될 때까지 반복한다. 유의해야할 점은 회전 초밥이기 때문에 rp가 인덱스 0을 가리키는 것이 가능하다는 것이다. 이것에 주의해서 구현하면된다.

2. 코드

import sys
from collections import defaultdict

n, d, k, c = map(int, sys.stdin.readline().rsplit())
c = str(c)
arr = [sys.stdin.readline().rstrip() for _ in range(n)]
kinds = defaultdict(int)
lp = 0
rp = lp + k - 1
count = 0
answer = 0
first = True

while lp < n:
    if first:
        for i in range(lp, rp+1):
            i %= n
            sushi = arr[i]
            if not kinds[sushi]: count += 1
            kinds[sushi] += 1
            first = False
        
        answer = max(answer, count) if kinds[c] else max(answer, count+1)
        continue

    first_sush = arr[lp]
    kinds[first_sush] -= 1
    if not kinds[first_sush]: count -= 1
    lp += 1
    rp = (rp + 1) % n
    
    if not kinds[arr[rp]]: count += 1
    kinds[arr[rp]] += 1
    answer = max(answer, count) if kinds[c] else max(answer, count+1)

print(answer)