[Python] 백준 / gold / 17294 : 오큰수

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

"크기가 N인 수열 A = A1, A2, ..., AN이 있을 때, 수열의 각 원소 Ai의 오큰수는 오른쪽에 있으면서 Ai보다 큰 수 중에서 가장 왼쪽에 있는 수를 의미한다. 그러한 수가 없는 경우에 오큰수는 -1이다."

나이브하게 해법을 생각한다면 이중 for 문을 사용한 O(N^2) time 의 해법을 생각할 수 있다. 근데 이 문제의 input 조건은 N (1 ≤ N ≤ 1,000,000) 이기 때문에, 그렇게 풀면 시간초과가 나게 된다.

O(log N) 이하의 시간복잡도를 가지는 자료구조 / 알고리즘을 떠올려보면, 이진탐색, 정렬, 해쉬 탐색, 다이나믹 프로그래밍, 힙.. 등이 있다.

이 문제의 출제 의도는 힙이었다. for 문으로 배열을 순차 탐색하면서 최소 힙에 원소들을 push 한다. 현재 넣으려는 원소가 힙의 top 보다 작다면 그냥 push 하고, 크다면 현재 원소보다 작은 원소들을 pop 해나가면 된다.

기억할 것 : heapq 에서 top 인덱스는 h[0] 으로 접근한다.

파이썬 코드

import heapq
import sys
N = int(sys.stdin.readline())
A = list(map(int, sys.stdin.readline().split()))
B = [-1] * N
h = []

for i in range(N):
    a = A[i]
    # 힙이 비어있으면 그냥 삽입
    if not h:
        heapq.heappush(h, (a, i))
        continue

    top = h[0]
    if a <= top[0]:
        heapq.heappush(h, (a, i))
    else:
        while h and a > h[0][0]:
            value, idx = heapq.heappop(h)
            B[idx] = a
        heapq.heappush(h, (a, i))

for b in B:
    print(b, end=' ')