[Python] 백준 15961 - 회전 초밥 문제 풀이

Overview

BOJ 15961번 회전 초밥 Python 문제 풀이
분류: Two Pointers (투 포인터)

---

문제 페이지

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

---

풀이 코드 1

from sys import stdin
from collections import defaultdict, deque


def main():
    def input():
        return stdin.readline()

    N, d, k, c = map(int, input().split())
    sushis = [int(input()) for _ in range(N)]

    window = deque()
    cnts = defaultdict(int)

    for i in range(k):
        window.append(sushis[i])
        cnts[sushis[i]] += 1

    res = 0
    cnt = len(cnts)
    if not cnts[c]:
        res = cnt + 1

    for start in range(N - 1):
        prev = window.popleft()
        cnts[prev] -= 1
        if cnts[prev] == 0:
            cnt -= 1
        window.append(sushis[(start + k) % N])
        cnts[window[-1]] += 1
        if cnts[window[-1]] == 1:
            cnt += 1
        if cnts[c] == 0:
            res = max(cnt + 1, res)
        else:
            res = max(cnt, res)

    print(res)


if __name__ == "__main__":
    main()

window에는 먹은 초밥을 저장하고 cnts 에는 먹은 초밥 가짓수를 저장한다. 투포인터로 먹은 초밥을 조정하면서 탐색하고, 가짓수가 최대가 될 때 결과값 res에 저장한다. 이때 window에 쿠폰 초밥이 들어있는지 여부에 따라 결과값을 조정한다.

---

풀이 코드 2

from sys import stdin
from collections import defaultdict


def main():
    def input():
        return stdin.readline().rstrip()

    N, d, k, c = map(int, input().split())
    sushis = [int(input()) for _ in range(N)]
    plates = defaultdict(int)
    plates[c] = 1

    cnt = 1
    for i in range(k):
        if plates[sushis[i]] == 0:
            cnt += 1
        plates[sushis[i]] += 1

    res = cnt
    for end in range(k, N + k - 1):
        plates[sushis[end - k]] -= 1
        if plates[sushis[end - k]] == 0:
            cnt -= 1

        plates[sushis[end % N]] += 1
        if plates[sushis[end % N]] == 1:
            cnt += 1

        res = max(cnt, res)

    print(res)


if __name__ == "__main__":
    main()

이전 코드를 줄이고 속도를 개선하였다.
쿠폰 번호 초밥은 이미 하나 먹었다고 설정하여 투포인터 탐색을 진행한다.