[BOJ] #2531 회전초밥 (Python, Sliding window, 시간초과 해결)

🌱 문제

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

 

🌱 풀이

0. 입력받기

input = sys.stdin.readline

n, d, k, c = map(int, input().split())
belt = []
for _ in range(n):
    belt.append(int(input()))

 

*처음 생각했던 방법 (시간 초과 발생)

최대 접시 수는 30,000개이니까 브루트 포스로 충분히 가능하다고 생각했습니다. 리스트를 돌면서 초밥의 종류만 세주면 되네!하고 아래와 같이 코드를 짰는데 시간초과가 발생하였습니다.

answer = 0

for i in range(n):
    belt.append(belt.popleft())
    answer = max(len(set(list(belt)[:k] + [c])), answer)

print(answer)

그 이유는 우선 for문이 $O(N)$의 시간 복잡도를 가지고 $list(belt)[:k]$에서 리스트로 변환 후 슬라이싱하는데에서 $O(k)$ 시간복잡도를 가집니다. 그리고 리스트를 set로 변환하는데에도 $O(k)$의 시간복잡도를 가지므로 최종적으로 for문 안에서 $O(k)$ 시간복잡도가 발생하는 것입니다. 따라서 위의 코드의 전체 시간복잡도는 $O(n * k)$가 됩니다.

 

1. 투포인터(슬라이딩 윈도우) 사용하기

위의 시간 초과를 해결하기 위해서는 투포인터(슬라이딩 윈도우) 알고리즘을 사용해야합니다.

window = defaultdict(int)

우선 defaultdict로 딕셔너리를 하나 만듭니다. 이 딕셔너리는 초밥의 종류를 세기 위한 딕셔러니입니다.

 

2. 윈도우 초기화

# Initialize the first window
for i in range(k):
    window[belt[i]] += 1

# Add the coupon sushi
window[c] += 1
answer = len(window)

첫번째 윈도우를 설정합니다. 첫번째 초밥부터 k개를 연속으로 먹고 쿠폰 초밥까지 먹었을 때의 초밥의 종류를 구합니다.

 

3. 투포인터 옮기기

# Slide the window across the belt
for i in range(1, n + 1):

    # Remove the sushi that goes out of the window
    window[belt[i - 1]] -= 1
    if window[belt[i - 1]] == 0:
        del window[belt[i - 1]]
        
    # Add the new sushi that enters the window
    window[belt[(i + k - 1) % n]] += 1
    
    # Update the maximum number of unique sushis
    answer = max(len(window), answer)

print(answer)

for문으로 belt 끝까지 돌면서 왼쪽 초밥 하나는 빼고 오른쪽 초밥 하나는 더해가면서 초밥의 종류를 세서 answer에는 최댓값이 들어가도록 합니다.

이렇게 코드를 짜게 되면 for문 안에서는 모든 코드의 시간복잡도가 $O(1)$이 되기 때문에 전체 코드는 $O(N)$의 시간복잡도를 가집니다.

 

🌱 정답 코드

# !/usr/bin/env python
# -*- coding: utf-8 -*-
# boj 2531 회전초밥

import sys
from collections import defaultdict

input = sys.stdin.readline

n, d, k, c = map(int, input().split())
belt = []
for _ in range(n):
    belt.append(int(input()))

current_window = defaultdict(int)

for i in range(k):
    current_window[belt[i]] += 1

current_window[c] += 1
answer = len(current_window)

for i in range(1, n):
    current_window[belt[i - 1]] -= 1
    if current_window[belt[i - 1]] == 0:
        del current_window[belt[i - 1]]
    current_window[belt[(i + k - 1) % n]] += 1
    answer = max(answer, len(current_window))

print(answer)