[17298] 오큰수 #python

백준 17298번 오큰수

내가 푼 코드 ⇒ 시간초과

num = int(input())
a = list(map(int, input().split(" ")))
result = []
for i in range(num):
    if i+1 == num:
        result.append("-1")
        break
    max_num = max(a[i+1:])
    if a[i] < max_num:
        for j in range(i+1, num):
            if a[j] > a[i]:
                result.append(str(a[j]))
                break
    else:
        result.append("-1")
print(" ".join(result))

정답

num = int(input())
a = list(map(int, input().split(" ")))
result = ["-1" for _ in range(num)]
stack = []
stack.append(0)
i = 1
for i in range(num):
    while stack and a[stack[-1]] < a[i]:
        result[stack[-1]] = str(a[i])
        stack.pop()

    stack.append(i)
    i += 1

print(" ".join(result))

num의 최댓값이 백만이므로, 이중포문을 돌면 시간초과가 발생한다.

stack을 사용해서 해결하는데, 작은 수의 인덱스를 stack에 push 하다가 큰 수를 만나면 큰 수보다 작은 stack의작은 수를 모두 pop한다. 이런 식으로 push와 pop을 반복하며 바로 다음 큰 수를 찾는다.