[백준17298번] 오큰수 / Python3

문제

크기가 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)을 공백으로 구분해 출력한다.

예제 입력 1

4
3 5 2 7

예제 출력 1

5 7 7 -1

예제 입력 2

4
9 5 4 8

예제 출력 2

-1 8 8 -1

나의 풀이

이 문제는 O(N^2)이 걸릴수도 있는 문제를 스택을 사용해 O(N)에 가깝게 만드는 문제였습니다.

왼쪽에서부터 값을 스택에 넣으면서 알고리즘을 진행해나가는데 스택에 값을 넣기 전에 스택에 가장 위에 있는 값이 현재의 값보다 작거나 스택이 빌때까지 pop연산을 해서 빼주고나서 현재의 값을 스택에 push를 해나가는 것입니다. 여기서 스택의 가장 위의 값이 pop이 되는 순간이 바로 해당 인덱스의 오큰수가 현재의 값이 된다는 의미가 됩니다. 이 과정을 그림으로 그리면 아래와 같습니다.

왼쪽부터 시작을 하게되고 처음에는 스택이 비어있으므로 3을 스택에 넣고 오른쪽 원소로 이동을 하게됩니다.

2의 경우 3보다 작기때문에 스택에 넣고 다음 원소로 이동합니다.

여기서 4를 스택에 넣기전에 스택의 가장 맨위의 값인 2와 비교해보니 2보다 크기때문에 pop을 함과 동시에 오큰수를 현재 값인 4로 업데이트 해줍니다. 이를 반복하여 3보다도 크기때문에 pop과 동시에 오큰수를 업데이트 합니다. 따라서 현재 3과 2의 오큰수는 4로 확정이 지어집니다. 이제 스택이 비었으니 4를 push하고 다음 원소로 넘어갑니다.

1은 4보다 작으므로 스택에 push를 하고 다음 원소로 넘어가게 됩니다.

여기서도 4때와 마찬가지로 1과 4가 모두 8보다 작으므로 pop과 동시에 오큰수를 8로 업데이트를 해줍니다. 그리고 8을 스택에 넣습니다. 마지막 수의 경우 오른쪽에 숫자가 존재하지 않기때문에 무조건 -1값이 됩니다. 따라서 여기서 반복을 멈추고 구한 오큰수를 출력해주면 문제를 해결하게 됩니다.

참고해야할 점은 스택에 원소를 집어넣을 때 해당 값의 인덱스를 함께 넣어주어야 스택에서 pop을 하면서 오큰수를 업데이트하는데에서 시간을 줄일 수 있습니다.

코드

# 17298번 오큰수
from collections import deque

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

oh_big = [-1] * n
stack = deque()

for i in range(n):
    while stack and (stack[-1][0] < seq[i]):
        tmp, idx = stack.pop()
        oh_big[idx] = seq[i]
    stack.append([seq[i], i])

print(*oh_big)