백준 17298번_오큰수 (Python)

문제

난이도: 골드4

크기가 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이다.

 

입력: 첫째 줄에 수열 A의 크기 N (1 ≤ N ≤ 1,000,000)이 주어진다. 둘째 줄에 수열 A의 원소 A1, A2, ..., AN (1 ≤ Ai ≤ 1,000,000)이 주어진다.

 

출력: 총 N개의 수 NGE(1), NGE(2), ..., NGE(N)을 공백으로 구분해 출력한다.

 

코드

from collections import deque

N = int(input())
A = list(map(int, input().split()))
NGE = [-1]*N
stack = deque([0])
for i in range(1, N):
    while len(stack) and A[stack[-1]]<A[i]:
        s = stack.pop()
        NGE[s] = A[i]
    stack.append(i)
for i in range(N):
    print(NGE[i], end = ' ')

 

해설

맨 앞의 원소부터 오른쪽으로 순서대로 검토한다. 이때 각 원소들은 자기보다 더 큰 수가 나오기 전까지 '대기 상태'에 있어야 한다. 여기서 조금 생각을 해보면 내림차순으로 된 스택을 사용할 수 있음을 알게 된다.

일단 대기 상태의 원소들을 대기 상태를 나타내는 배열 맨 끝에 계속 집어넣는다고 하자. 만약에 오름차순이 되는 특정 구간이 배열 안에 존재한다면, 왼쪽의 더 작은 원소보다 이미 더 큰 항목이 검토된 상황이므로 대기 상태가 아니어야 한다. 따라서 하나의 새 원소를 검토할 때마다 앞의 원소들의 대기 상태를 해제해주는 과정을 성실히 이행해주면 배열은 내림차순으로 유지될 것이다.

이때 대기 상태를 해제해주는 작업을 어떻게 이행할지 살펴본다. 그런데, 앞에서 대기상태를 해제해주기만 해도 배열이 내림차순으로 되어 있기 때문에 새 원소를 집어넣어 내림차순이 되게 하기 위해서는 맨 끝에서부터 새 원소보다 작은 원소들을 꺼낸 뒤, 새 원소를 맨 끝에 넣어도 내림차순이 될 때 집어넣으면 된다. 맨 끝으로 넣어서 맨 끝으로 빼내니, 이것은 스택으로 구현할 수 있음을 알 수 있다.

따라서 스택에 새 원소를 추가하려면, 1. 새 원소보다 작은 원소들을 모조리 다 꺼집어내고, 2. 이때까지 그 원소들보다 큰 원소가 없었으니 대기 상태일 것이므로 새 원소를 그들의 NGE로 정한 후, 3. 새 원소 이상의 원소만 남았을 때 스택에 추가한다.