[Python] 백준 - 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이다.

내 코드

1차 시도

import sys
N = int(sys.stdin.readline())
nge = list(map(int, sys.stdin.readline().split()))
ans = []
for i in range(N):
    ans = [x for x in nge[i:] if x > nge[i]]
    if ans:
        print(ans[0], end=' ')
    else:
        print(-1, end=' ')

시간 초과

 

2차시도

import sys
N = int(sys.stdin.readline())
nge = list(map(int, sys.stdin.readline().split()))
# N = 4
# nge = 9, 5, 4, 8
ans = []
for i in range(N):
    flag = True
    for j in nge[i:]:
        if nge[i] < j:
            print(j, end=' ')
            flag = False
            break
    if flag:
        print(-1, end=' ')
    # nge = nge[i:]

시간 초과

 

찾아보니 입력이 최대 1,000,000까지여서 완전 탐색으로 풀면 O(N^2)의 시간이 걸리기 때문에 시간 초과가 발생한다고 한다.

솔루션 검색 후 다시 작성한 코드

import sys

N = int(sys.stdin.readline())
num = list(map(int, sys.stdin.readline().split()))
NGE = [-1]*N
stack = []

for i in range(N):
    while stack and num[stack[-1]] < num[i]:
        NGE[stack.pop()] = num[i]
    stack.append(i)
print(*NGE)

N개의 -1로 구성되어 있는 만들고 오큰수가 존재하는 인덱스만 값을 추가한다.
stack은 인덱스를 저장하는 공간으로 사용한다.

여태 푼 문제 중 가장 시간이 오래 걸렸다.