[백준][17298] 오큰수

• Problem Link : https://www.acmicpc.net/problem/17298

문제

자신보다 오른쪽에 있으면서 큰 수 중에서 가장 왼쪽이 있는 수를 오큰수라고 한다. 없는 경우에는 -1이라고 한다. 크기가 N인 수열이 A가 주어질 때 각각의 오큰수를 구해보아라.

입력

• 첫째 줄, 수열의 A의 크기 N이 주어진다.
• 둘째 줄, 수열 A가 주어진다.

출력

• N개의 오큰수를 출력한다.

풀이

수열 A를 오른쪽에서 왼쪽 방향으로 진행하면서 숫자를 스택에 push하는데 들어오는 숫자보다 작은 숫자들을 모두 pop한 후에 push를 해준다. 이러면 그 다음으로 왼쪽 원소 차례가 되면 어차피 오른쪽 첫번째 원소는 방금 pop을 모두 해준 후 push가 된 그 원소이기 때문에 정상적인 오큰수 계산이 가능하다.

• A = [3, 5, 2, 7]를 예로 들어본다면 아래와 같다.

1. 7의 경우 스택이 비었으므로 답에 -1 추가 후 스택에 7 push
2. 2의 경우, 스택의 맨 위의 7보다 작으므로 오큰수로 답에 7을 추가 후 스택에 2 push
3. 5의 경우, 스택의 맨 위의 2가 자신보다 작으므로 pop
4. 5의 경우, 그 다음 스캑의 맨 위의 7이 자신보다 크므로 오큰수로 답에 7을 추가 후 스택에 자신의 5을 push
5. 3의 경우, 스택의 맨 위의 5가 자신보다 크므로 답에 5을 추가 후 3을 push
6. 모든 숫자를 돌았으므로 답은? [-1, 7, 7, 5]
7. 왼쪽에서 돌았으므로 reverse 시켜주면? [5, 7, 7, -1]로 원하는 답이 나온다.

위 과정처럼 크기를 비교하면서 스택에 자신보다 작은수가 올라와있다면 수를 모두 pop 해준 후 올라오는 수를 답 리스트에 추가, 만약 큐에 아무것도 없다면 자신보다 큰 수가 오른쪽에 없었다는 것이므로 -1을 답 리스트에 추가를 반복하면 된다.

코드

import sys
from collections import deque

input = sys.stdin.readline

# 오른쪽에서 왼쪽으로 push하면서 진행
# 들어오는 숫자가 스택의 맨 위 숫자보다 크다면 pop

n = int(input())
a = list(map(int, input().split()))

q = deque()
ans = []

for i in range(n - 1, -1, -1):
    while q and q[-1] <= a[i]:
        q.pop()
    if not q:
        ans.append(-1)
    else:
        ans.append(q[-1])
    q.append(a[i])
ans.reverse()
print(*ans)