백준 2531 회전 초밥 (파이썬)

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

 

첫번째 풀이 (pypy3 2772ms, python3 시간초과)

import sys
def solve():
    left = 0
    right = k
    answer = 1
    while left < N:
        eat = set()
        # 초밥을 k개 만큼 먹기
        for i in range(left, right):
            eat.add(arr[i%N])
        # c가 없으면 1개 추가하기
        eat.add(c)
        answer = max(answer, len(eat))
        # 1칸씩 밀기
        left += 1
        right += 1

    print(answer)
    return

if __name__=="__main__":
    input = sys.stdin.readline
    # N: 접시의 수
    # d: 초밥의 가짓수
    # k: 연속해서 먹는 접시의 수
    # c: 쿠폰 번호
    N, d, k, c = map(int, input().split())
    arr = []
    susi = [0]*(d+1)
    for i in range(N):
        num = int(input().strip())
        arr.append(num)
    solve()

K개의 연속된 초밥들 중에서 set으로 중복을 제거하여 유니크한 초밥의 번호만 체크를 했습니다

원형 리스트는 i%N 으로 right가 N을 넘어가도 다시 첫 초밥부터 시작하게 만들었습니다.

그리고 left가 N을 넘어가면 종료를 시켜주었습니다.

하지만 시간복잡도가 O(N*K)라서 최대 9억번을 연산하는 제 코드는 python3 로는 해결할 수 없었습니다.

 

두번째 코드(python3, pypy3 모두 통과)

import sys

def solve():
    left = 0
    right = 0
    dict = {}
    result = 0
    # k만큼 먹기
    while right < k:
        if arr[right] not in dict.keys():
            dict[arr[right]] = 1
        else:
            dict[arr[right]] += 1
        right += 1

    # c는 반드시 추가
    if c not in dict.keys():
        dict[c] = 1
    else:
        dict[c] += 1

    # 슬라이딩 윈도우
    while left < n:
        result = max(result, len(dict))
        # 1. 맨 왼쪽 초밥 제거
        dict[arr[left]] -= 1
        # 삭제된 왼쪽 초밥이 0 이면 dict 삭제
        if dict[arr[left]] == 0:
            del dict[arr[left]]
        if arr[right % n] not in dict.keys():
            dict[arr[right % n]] = 1
        else:
            dict[arr[right % n]] += 1
        left += 1
        right += 1

    print(result)
    return

if __name__ == "__main__":
    n, d, k, c = map(int, sys.stdin.readline().split())
    arr = [int(sys.stdin.readline()) for _ in range(n)]
    solve()

검색을 해본 결과 딕셔너리로 푸는 방법을 보고 따라서 해봤습니다. 

dict 는 반복문을 사용하지 않고 그 안에서 개수가 0인지만 체크 하기 때문에 빠르게 작동하는 것을 확인할 수 있었습니다.

 

세번째 코드 (defaultdict 사용)

import sys
from collections import defaultdict

def solve():
    left = 0
    right = 0
    dict = defaultdict(int)
    result = 0
    # k만큼 먹기
    while right < k:
        dict[arr[right]] += 1
        right += 1

    # c는 반드시 추가
    dict[c] += 1

    # 슬라이딩 윈도우
    while left < n:
        result = max(result, len(dict))
        # 1. 맨 왼쪽 초밥 제거
        dict[arr[left]] -= 1
        # 삭제된 왼쪽 초밥이 0 이면 dict 삭제
        if dict[arr[left]] == 0:
            del dict[arr[left]]
        dict[arr[right % n]] += 1
        left += 1
        right += 1

    print(result)
    return

if __name__ == "__main__":
    n, d, k, c = map(int, sys.stdin.readline().split())
    arr = [int(sys.stdin.readline()) for _ in range(n)]
    solve()

 

2번째 코드는 매번 key가 딕셔너리에 있는지 확인 후 value를 넣어주는 불편함이 있습니다.

defaultdict를 이용해서 매번 키 확인을 하지 않고 키가 없으면 바로 value를 0으로 갖는 key를 생성하게 코드를 구현해봤습니다. 

아래는 제가 직접 디버깅 할 때 사용한 반례들 입니다.

8 30 4 30 
7
7
3
7
7
8
7
7
정답 4

2 2 2 2
1
1
정답 2

6 6 6 6
1
2
3
4
5
6
정답 6

21 435 20 408
155
257
144
356
162
173
409
206
60
101
254
380
122
18
14
372
206
404
249
272
169
정답 21

8 30 4 30
7
9
7
30
2
9
7
25
정답 5

 

결과 및 정리 

처음에 문제를 제대로 이해하지 못해서 시간만 날리구 검색으로 제대로 이해했습니다.

서로 다른 초밥의 종류를 4개 먹는 것이 아니라 연속으로 4개만 먹으면 되는 것이었습니다.

처음에 이렇게 이해는 했지만 이렇게 쉬울 수 없다고 생각해 혼자서 문제를 한번 더 꼬아서 생각한 것이 잘못이었습니다.

defaultdict를 처음 써봤는데 너무 편리해서 이제부터 기본 딕셔너리대신 defaultdict를 불러와 사용할 것 같습니다.