백준 17928번 오큰수

백준 17928번: 오큰수

아이디어

수열의 크기가 1,000,000이다. O(NlogN)의 시간복잡도를 가진 알고리즘을 설계하면 풀 수 있을 것 같다. 그런데 이 문제는 O(N)의 시간복잡도로 풀 수 있다.

1. 입력받은 수열을 역순으로 집어넣은 스택 A, 빈 스택 B를 만든다.
2. 스택 B가 비어있으면 A에서 pop해서 스택 B에 push한다.
3. 스택 B에 원소가 있다면, 스택 A의 최상단에 위치한 원소와 스택 B의 최상단에 위치한 원소를 비교한다.
4. 스택 A의 최상단에 위치한 원소가 스택 B의 최상단에 위치한 원소보다 크다면, NGE(그 원소의 인덱스)를 스택 A의 최상단에 위치한 원소로 업데이트하고 pop한다.
5. 스택 B의 최상단에 위치한 원소가 스택 A의 최상단에 위치한 원소보다 값이 커지거나, 스택 B에 원소가 없어질 때까지 반복한다.
6. 스택 A에서 pop해서 스택 B에 push한다.

코드

import sys
input = sys.stdin.readline

N = int(input())
st = list(map(int, input().split()))

targets = []
ans = [-1] * N

st.reverse()
idx = 0
targets.append((idx, st.pop()))

while st:
    t = st.pop()
    idx += 1

    while targets and targets[-1][1] < t:
        a, b = targets.pop()
        ans[a] = t

    targets.append((idx, t))

for i in ans:
    print(i, end=' ')

여담

문제풀기보다 그림그려서 설명하는게 어려움