[Python] 백준 / gold / 2493 : 탑

문제 링크 : https://www.acmicpc.net/problem/2493

내가 스택에 대한 이해도가 부족했다고 느끼게 해준 문제. N = 5x10^5 이기 때문에 일단 O(N^2) 의 알고리즘으로는 택도 없는건 당연하다.

스택, 큐, 힙, 이분탐색, 해쉬 ... 여러 풀이 방법을 떠올려봤다. 근데 이런 개념들을 알고 있으면 뭐해, 제대로 써먹질 않는데. 바보같게도 스택으로 정답에 근접한 풀이를 생각해놓고도 "아 이럼 시간초과 날것같은데" 생각하고 멀리 떠나버렸다.

처음에 떠올린 스택으로 바로 풀었으면 30분 이상 단축했을꺼다. 그래도 이 문제 덕분에 스택에 대해서 다시 짚어봤고.. 이해도를 높일 수 있었던 좋은 기회라고 생각해야겠다.

기억할 것 : 스택에서는 pop 이 일어나기 때문에 N 이 변한다. 시간복잡도를 계산할 때 이거 띵킹하자.

파이썬 코드

import sys
N = int(sys.stdin.readline())
tower = list(map(int, sys.stdin.readline().split()))
answer = [0]
stack = [(tower[0], 1)]

for i in range(1, N):
    t = tower[i]
    if not stack:
        stack.append((t, i+1))
        answer.append(0)
        continue

    if t > stack[-1][0]:
        while stack and t > stack[-1][0]:
            stack.pop()

        if not stack:
            answer.append(0)
        else:
            answer.append(stack[-1][1])
        stack.append((t, i+1))
    else:
        answer.append(stack[-1][1])
        stack.append((t, i+1))

for a in answer:
    print(a, end=' ')