백준 - 회전초밥 / Silver 1 / 2531번 / Python

문제 📋

백준 - 회전초밥

처음 풀이 📝

import sys
from collections import deque

n, d, k, c = map(int, sys.stdin.readline().split())

sushi = [int(sys.stdin.readline()) for _ in range(n)]


def solution(n, d, k, c, sushi):
    answer = []
    choice = deque()

    for i in range(k):
        choice.append(sushi[i])
    kind = set(choice)
    kind.add(c)
    answer.append(len(kind))

    for i in range(1, n):
        choice.popleft()

        new_sushi = sushi[(i+k-1) % n]
        choice.append(new_sushi)

        kind = set(choice)
        kind.add(c)
        answer.append(len(kind))

    return max(answer)


print(solution(n, d, k, c, sushi))

원형으로 된 벨트 위의 초밥을 k개 연속으로 먹었을 때 가장 많은 종류의 수를 구하는 문제이다.
쿠폰이라는 변수가 있지만 쿠폰으로 제공되는 초밥은 이미 먹은 것으로 치면 된다.

처음 접근한 방식은 deque에 k개의 초밥들을 관리하면서 set 변환을 통해 종류 수를 answer 리스트에 넣어주는 방식이었다.

이 방식으로도 모든 테스트케이스를 통과했지만
반복문내에서 set() 함수를 호출하기 때문에 시간복잡도가 O(kn)으로 불리한 측면이 있었다.
접시의 수인 n이 최대 30,000 k가 최대 3,000이기 때문에 최대 연산횟수가 9천만 가까이 되어 시간제한은 통과했지만 매우 느렸다.

개선된 풀이 📝

import sys
from collections import defaultdict

n, d, k, c = map(int, sys.stdin.readline().split())

sushi = [int(sys.stdin.readline()) for _ in range(n)]


def solution(n, d, k, c, sushi):
    choice = defaultdict(int)  # 초밥 선택정보를 딕셔너리로 관리
    choice[c] += 1  # 쿠폰으로 받는 초밥이 이미 먹었다고 가정

    for i in range(k):  # 초기 k개 초기화
        choice[sushi[i]] += 1
    answer = len(choice)

    for i in range(n):
        start = i
        next_sushi = (i+k) % n

        choice[sushi[start]] -= 1  # 시점에 있는 초밥 선택을 되돌림

        if choice[sushi[start]] == 0:  # 해당 초밥을 1개도 먹지 않았다면
            del choice[sushi[start]]  # 종류에 카운팅 되지 않도록 배제

        choice[sushi[next_sushi]] += 1  # 새로운 초밥 선택에 추가

        answer = max(answer, len(choice))  # 초밥 종류 수의 최대값 갱신
    return answer


print(solution(n, d, k, c, sushi))

따라서 더 좋은 풀이를 고민하였고 위와 같이 풀었다.
먼저 굳이 집합을 사용하지 않고 딕셔너리를 통해 초밥 선택정보를 관리해주었다.
또한 최대값을 하나로 관리하고 갱신하는 방법을 통해 시간복잡도를 O(n)으로 크게 줄일 수 있었다.