[BOJ] 17298 : 오큰수 (Python3)

출처: https://www.acmicpc.net/problem/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)을 공백으로 구분해 출력한다.

소스코드

import sys
input = sys.stdin.readline

n = int(input())
 # 수열을 받아두는 리스트
inputs = list(map(int, input().split()))
 # 오큰수를 찾지 못한 inputs의 index를 모아두는 리스트
stack = [0]
 # 출력할 오큰수를 모아두는 리스트. 미리 -1로 설정해두고 오큰수를 찾은 값만 바꿔준다.
outputs = [-1 for _ in range(n)]

for i in range(1, n):
     # 오큰수를 찾기 위한 반복문
    while stack and inputs[stack[-1]] < inputs[i]:
        outputs[stack[-1]] = inputs[i]
        stack.pop()
    stack.append(i)
    i += 1

for i in outputs:
    print(i, end=' ')

풀이

쉽게 생각하면 이중 반복문을 쓰게되는데 시간복잡도가 O(N^2), N의 범위가 1 ≤ N ≤ 1,000,000 이므로 시간초과가 뜬다.
때문에 스택을 사용해서 풀어야 한다.
주의할 점은 스택에 결과값이 아니라 아직 오큰수를 찾지 못한 inputs의 index를 넣어둬야 한다.

그래서 inputs[i] 값이 inputs[stack[-1]] (오큰수를 찾지 못한 수 중 가장 위에 있는 것) 보다 크다면, 즉 오큰수를 찾았다면
outputs의 i 자리에 있는 값을 inputs에 있는 값으로 바꿔준다. 그 인덱스는 오큰수를 찾았으므로 stack에서 pop해준다.
stack에 있는 다른 인덱스 또한 inputs[i] 보다 작다면 오큰수를 찾은 것이므로 반복해준다.

오큰수를 다 찾았거나 없다면 다음 인덱스로 넘어가야하므로 stack에 i를 추가해주고 i += 1을 해준다.