백준 17298 - 오큰수 [파이썬]

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

17298번: 오큰수

풀이 과정

문제를 다 읽으면 수열을 다 입력받은 상태에서 Ai의 다음 원소부터 끝까지 검색하면서 Ai보다 큰 수가 보이면 바로 출력해주는 방식이 떠오르지만 이 방법은 O(N^2) 방법이고, A의 크기가 100만이므로 시간 제한 1초 내로 해결할 수 없다.

 

스택을 이용한 더 효율적인 방법으로 풀어보자. 스택을 하나 생성하고, 오큰수가 담겨있는 배열 NGE를 만들고 -1로 초기화 하자.

 

입력 배열 A의 첫 번째 원소부터 살펴보면서 오큰수를 아직 못 찾은 수의 index를 스택에 저장하면서, 다음 index j를 확인했을 때 스택탑의 index i와 비교해 Ai < Aj이면 NGE[i] = Aj로 두고 스택에서 원소를 꺼내는 과정을 Ai < Aj가 아닐 때까지 반복해주자. 이러면 A의 원소에 대해 존재하는 모든 오큰수가 구해진다.

 

왜 이러한 방법이 성립할까? 스택에서 더 작은 수의 index일수록 위에 담겨있기 때문이다. 더 큰 수의 index는 i의 오큰수가 되어서 스택에 들어오지 않고 기존 수보다 더 작은 수의 index만 스택에 쌓여간다.

코드

import sys

N = int(input())
A = list(map(int, input().split()))
NGE = [-1 for x in range(N)]
stack = []

for j in range(N):
    if len(stack) == 0:
        stack.append(j)
        continue
    i = stack[-1]
    while A[i] < A[j]:
        NGE[i] = A[j]
        stack.pop()
        if len(stack) == 0: break
        else: i = stack[-1]
    stack.append(j)

for i in range(N): print(NGE[i], end=' ')