[Python] 백준 17298번 - 오큰수 (골드 4)

혼자 힘으로 풀었는가? O

알고리즘 분류
 - 스택

 

 

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

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이다.

입력

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

출력

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

 

---

 

 

어찌 풀어야하나.. 싶어서 알고리즘 분류를 봤는데 스택이었다.

 

그래서 어찌저찌 스택으로 푸는 법을 찾았다.

 

1. 처음부터 마지막까지 순회.

2. 스택이 비어있으면 현재 위치와 값을 스택에 넣는다.

3. 현재 위치와 스택의 최상단 비교

3-1. 현재 위치의 값 < 스택 최상단 이면 현재위치 스택에 추가

3-2. 현재 위치의 값 > 스택 최상단 이면 스택 최상단의 위치에 현재 위치의 값을 넣는다.

4. 스택이 비거나, 3-1 기준에 걸릴 때까지 반복

5. 2~4 반복

 

import sys
input = sys.stdin.readline

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

result = [-1] * n
stack = []

for i in range(n):
    if len(stack) == 0: # 비어있으면 넣기
        stack.append((data[i], i))
    else:
        while stack:
            top, top_idx = stack.pop()
            if data[i] > top: # 현재값보다 최상단이 크면 최상단의 위치에 현재 값 넣기
                result[top_idx] = data[i]
            else: # 현재값이 최상단보다 작으면 다시 원복
                stack.append((top, top_idx))
                break
        stack.append((data[i], i)) # 현재 위치 스택에 추가.
            
for i in range(n):
    print(result[i], end=' ')