[BaekJoon] 17298 오큰수 (Gold 4) - Python

문제

 

크기가 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

 

---

코드

 

import sys
input = sys.stdin.readline

N = int(input()) # N입력
num = list(map(int, input().rstrip().split())) # 숫자들 입력
result = [-1 for _ in range(N)] # 결과값을 저장할 리스트 생성
stack = [] # 사용할 스택 선언

for i in range(N): # N만큼 반복
    while stack and (stack[-1][0] < num[i]): # 스택이 비어있지 않고, 스택 맨 위의 값이 현재의 값보다 작은 경우
        tmp, idx = stack.pop() # 스택 맨 위의 값을 pop하여 각각 tmp와 idx에 저장
        result[idx] = num[i] # 결과값 리스트의 idx인덱스 값을 현재 값으로 변경
    stack.append([num[i], i]) # 스택에 num[i]와 i값 추가

for i in result: print(i, end = ' ') # 출력

---

실행 화면

 

---

채점 결과