파이썬) 회전초밥 - BOJ

https://www.acmicpc.net/problem/2531

2531번: 회전 초밥

 

이 문제는 꽤나 까다로웠다.

문제 이해와 풀이에는 무척 쉬웠지만 시간초과가 많이 발생하는 문제였다.

 

우선 처음 제출한 코드이다.

 

from collections import deque
import sys
input = sys.stdin.readline
n, d, k, c = map(int, input().split())

arr = []

for _ in range(n):
    arr.append(int(input()))

arr = arr

answer = 0

q = deque(arr)

j = 0

while j < n:
    s = set()
    is_True = True
    for i in range(0, k):
        num = q[i]
        s.add(num)
        if num == c:
            is_True = False
    cnt = len(s)
    # break

    if is_True:
        cnt += 1

    answer = max(answer, cnt)
    q.rotate(1)
    # print(cnt)
    # print(q)
    # break
    j += 1
print(answer)

 

아쉽지만 시간초과가 발생했다.

rotate때문인가 싶어서 다른 방법으로도 해봤지만 역시나 시간초과가 발생했다.

결국 다른 사람의 코드를 살펴보았다.

import sys
N, d, k, c = map(int, input().split()) # N: 접시의 수 / d: 초밥의 가짓수 / k: 연속해서 먹는 접시수 / c: 쿠폰 번호

# input = sys.stdin.readline
cp1, cp2 = 0, k-1
sushi = []
for _ in range(N):
    sushi.append(int(input()))
check = set()

while cp1 < N:
    if cp2 >= N:
        cp2 -= N
    if cp2 < cp1:
        plates = sushi[cp1:] + sushi[:cp2+1]
    else:
        plates = sushi[cp1:cp2+1]
    cases = set(plates)
    if c not in cases:
        cases.add(c)
    if len(check) < len(cases):
        check = cases
    cp1 += 1
    cp2 += 1
print(len(check))

위 코드는 시간초과가 발생하지 않는다.

plates라는 리스트에 슬라이싱하는데 더 많은 시간이 잡아먹힐 것 같지만 그렇지는 않은것 같다.

 

위 코드의 해설은 다음과 같다

 

cp1,cp2는 각각 초밥 인덱스의 i번째와 i+k번째이다.

cp1이 n보다 크거나 같다면 while문을 종료한다

cp2가 n보다 크거나 같다면 다시 cp2를 0번부터 돌린다.

cp2 < cp1 또한 마찬가지로 순열배열을 위해 넣은 코드이다.

 

슬라이싱으로 얻은 리스트를 set을 통해 중복을 제거해 cases라는 변수에 넣었다.

cases에 c가 있는지 체크, 없다면 set에 추가로 c를 넣어준다.

check라는 최종적으로 가장 많은 가짓수를 담은 변수에 len(cases)와 check중 큰 값을 넣는다.

cp1,cp2를 +1만큼 증가시킨다.

 

결국 내가 작성한 코드와 별반 다르지 않지만 저 코드는 시간초과가 발생하지 않았으니 내가 생각한게 정답이 아닌 모양이다.

어디서 시간 초과가 발생했는지 찾아야겠다.