백준 | 카드2

백준 카드2 문제 풀이이다.

아래의 두 행위를 카드가 1장 남을 때까지 반복할 때, 그 1장의 번호를 출력하는 문제이다.

[1] 제일 위에 있는 카드를 바닥에 버린다.
[2] 뒤이어 제일 위에 있는 카드를 제일 아래로 옮긴다.

아래는 시간 복잡도를 고려하지 않고, 문제에서 시키는 대로 구현한 코드이다.

import sys


n = int(sys.stdin.readline())

cards = [c for c in range(n, 0, -1)]
while True:
    if len(cards) == 1:
        break

    # 마지막 카드 바닥에 버리기
    cards.pop()
    # 마지막 카드 제일 아래로 옮기기
    card = cards.pop()
    cards.insert(0, card)

print(cards[0])

당연히(?) 시간 초과가 뜨는데 그 이유는 다음과 같다.

먼저 while 루프가 카드가 1장 남을 때까지 반복되므로 카드 수 n에 비례하여 O(n)이라는 시간이 소요될 것이다.

그리고 리스트의 마지막 원소를 제거하는 pop()은 O(1)의 시간 복잡도를 갖지만, insert()는 O(n)의 시간 복잡도를 갖는다.

따라서 위 코드는 전체적으로 O(n2)의 시간 복잡도를 갖게 되고, n의 범위가 1 이상 500,000 이하이므로 2초라는 시간 제한에 대해 시간 초과가 발생하게 된다.

여기서 시간 복잡도를 줄일 수 있는 방법은 뭘까?

카드가 1장 남을 때까지 주어진 행위를 반복해야 하는 부분을 최적화할 수도 있겠지만..
주어진 행위는 결국 자료 구조의 양쪽 끝에서 원소를 삽입하거나 삭제하는 것이므로 양쪽 끝에서의 삽입, 삭제 연산의 시간 복잡도가 모두 O(1)인 collections.deque를 활용할 수 있다!

이를 활용하면 while 루프에 의한 O(n) 말고, 연산에 의한 시간 복잡도는 O(1)으로 처리 가능하므로 전체적으로 O(n)의 시간 복잡도를 갖는 코드를 짤 수 있다.

코드(정답)는 다음과 같다.

import sys
from collections import deque

n = int(sys.stdin.readline())

cards = deque([c for c in range(1, n + 1)])
while True:
    if len(cards) == 1:
        break

    # 마지막 카드 바닥에 버리기
    cards.popleft()
    # 마지막 카드 제일 아래로 옮기기
    card = cards.popleft()
    cards.append(card)

print(cards[0])