[백준/Python] 17298번: 오큰수

문제

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

크기가 N인 수열 A = A1, A2, ..., AN이 있다. 수열의 각 원소 Ai에 대해서 오큰수 NGE(i)를 구하려고 한다. Ai의 오큰수는 오른쪽에 있으면서 Ai보다 큰 수 중에서 가장 왼쪽에 있는 수를 의미한다. 그러한 수가 없는 경우에 오큰수는 -1이다.
예를 들어, A = [3, 5, 2, 7]인 경우 NGE(1) = 5, NGE(2) = 7, NGE(3) = 7, NGE(4) = -1이다. A = [9, 5, 4, 8]인 경우에는 NGE(1) = -1, NGE(2) = 8, NGE(3) = 8, NGE(4) = -1이다.

풀이과정

N의 최대 크기가 1,000,000이므로, 단순히 이중 for문으로 오른쪽 값을 확인하면 시간 초과가 발생한다. 따라서 한 번의 순회로 오큰수를 찾기 위해 스택을 활용한다.

핵심 아이디어는 다음과 같다.

1. 스택에는 아직 오큰수를 찾지 못한 인덱스를 저장한다.
2. i번째 원소 lst[i]를 보면서,
   스택 top에 있는 인덱스들의 값이 lst[i]보다 작다면
   해당 값들의 오큰수는 lst[i]이므로 answer에 기록하고 스택에서 제거한다.
3. i는 아직 오큰수를 찾지 못한 상태이므로 스택에 push한다.

이 과정을 전체 배열에 대해 수행하면 각 원소의 오큰수를 O(N)에 구할 수 있다.

풀이

n = int(input())
lst = list(map(int, input().split()))
answer = [-1] * n

stack = [0] # 스택에는 아직 오큰수를 찾지 못한 수의 인덱스가 들어있다!
for i in range(1, n):
    while stack and lst[stack[-1]] < lst[i]:
        idx = stack.pop()
        answer[idx] = lst[i]
    stack.append(i)

print(*answer)