알고리즘/BAEKJOON (Python)

[ 백준 / G4 / Pypy3 ] 15961번 - 회전초밥

Heeto 2023. 5. 11. 10:28
링크 : https://www.acmicpc.net/problem/15961
 

15961번: 회전 초밥

첫 번째 줄에는 회전 초밥 벨트에 놓인 접시의 수 N, 초밥의 가짓수 d, 연속해서 먹는 접시의 수 k, 쿠폰 번호 c가 각각 하나의 빈 칸을 사이에 두고 주어진다. 단, 2 ≤ N ≤ 3,000,000, 2 ≤ d ≤ 3,000, 2

www.acmicpc.net

 

코드


from collections import defaultdict

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

for start in range(n+k):
    idx = start if start < n else start - n
    if kind[sushi[idx]] == 0:
        res += 1
    kind[sushi[idx]] += 1
    if start >= k:
        kind[sushi[start-k]] -= 1
        if kind[sushi[start-k]] == 0:
            res -= 1
    answer = max(answer,res + int(not kind[c]))

print(answer)

풀이


" 투 포인터 "

풀이


구현은 간단하지만 시간복잡도를 줄이는데에 집중해야 하는 문제이다.

 

먼저 그리디하게 인덱스를 선택해서 정답을 구할 방법은 없기 때문에 완전탐색으로 가야 한다. 결국 첫 번째 인덱스부터 N, 그리고 회전이 가능하기 때문에 N+K-1까지 탐색할 수 밖에 없다. 근데 이 모든 반복마다 슬라이싱으로 부분배열을 구하면 편하긴 하지만 시간복잡도가 크게 증가하여 통과하지 못한다.

 

이런 이유 때문에 투 포인터를 사용해야 한다. 저장할 배열의 크기를 k개로 고정하고 탐색을 이어나가면서 k개가 넘어가면 뒤에 원소를 추가해주면서 앞의 원소를 제거해주는 것이다.

그리고 내가 현재 저장한 배열을 진짜 배열을 만들면 공간복잡도와 시간복잡도가 늘어날 것 같아서 진짜 배열을 만들지 않고 딕셔너리를 생성했다. 딕셔너리에 초밥의 종류를 키값으로 설정하고 값이 0이면 현재 배열에 없는 것이고 값이 1이상이면 현재 배열에 있는 것이다.

if start >= k:
    kind[sushi[start-k]] -= 1
    if kind[sushi[start-k]] == 0:
        res -= 1

맨 처음에 배열을 미리 k개로 만들고 시작해도 좋으나 나는 0번 인덱스부터 시작하는 방법을 선택했다.

그래서 배열의 길이가 k가 넘을 때부터 맨 앞의 원소를 빼주기 시작했다.

맨 앞의 원소를 빼면서 만약에 배열에 그 원소가 아예 사라진다면 res에서 1을 빼주었다.

 

    if kind[sushi[idx]] == 0:
        res += 1
    kind[sushi[idx]] += 1

만약 배열에 해당 초밥이 없다면 res에 1을 더해줌으로써 탐색중인 초밥의 종류를 증가시켜준다.

 

answer = max(answer,res + int(not kind[c]))

현재 탐색중인 배열에 쿠폰 초밥이 포함되어 있지 않다면 1을 더해서 정답을 갱신해준다.