[백준] 2531번 회전 초밥 파이썬 코드

코드1

import sys
input = sys.stdin.readline

N,d,k,c = map(int, input().split())

res = 0
sushi = [int(input()) for _ in range(N)]

for i in range(N):
  if i+k+1 <= N:
    eatdish = sushi[i:i+k] # [1,2,3,4,5]에서 슬라이싱할 때 1,2,3,4가 되는 경우
  else:
    eatdish = sushi[i:]+sushi[:k-(N-i)] # 위 경우에서 3,4,5,1로 되는 경우

  eatdish = set(eatdish) # 중복 제거를 위해 set 사용
  eatdish.add(c)
  res = max(res,len(eatdish))
print(res)

 데이터 수가 많았지만, 슬라이싱을 이용하면 시간제한에 걸리지 않을 것 같아서 위 방법을 선택했다. 회전초밥의 특성 상, 캐이스를 분류하여 해결했다. 그런데 나머지 정리를 이용하면 더 쉽게 해결할 수 있는 방법이 있어서 감탄했다...

 

코드2

import sys
input = sys.stdin.readline

N,d,k,c = map(int, input().split())
sushi = [int(input()) for _ in range(N)]
res = 0

for i in range(N):
  visited = [0]*(d+1)
  cnt = 0
  for j in range(i, i+k):
    dish = sushi[j%N] # 나머지 정리 이용
    if not visited[dish]:
      visited[dish] = 1
      cnt += 1
  if not visited[c]:
    cnt += 1
  res = max(res, cnt)
  
print(res)

그런데 이 코드는 pypy에서만 통과 가능하다. 데이터 수가 너무 많아서 그런 듯하다.