백준 17298번 "오큰수"

문제

백준 17298번 오큰수

풀이

문제 자체는 어렵지 않지만 시간을 줄여야해서 신경쓰이던 문제.
처음엔 이중 for 반복을 이용하여 풀었으나 계속해서 시간초과가 뜨길래 찾아보니 스택을 이용하여 풀어야 한다더라...

다른 스택 관련 문제와 다른 점은 인덱스를 스택에 넣어야 하는 점이다.
해당 인덱스의 수보다 큰 최초의 오른쪽 수를 찾을 때마다 그 수를 배열에 저장하여
O(n^2)의 시간복잡도를 O(n)까지 줄일 수 있게 된다.

Python 코드

import sys
input = sys.stdin.readline

n = int(input())
arr = list(map(int, input().split()))
answer = [-1 for _ in range(n)]
stack = []

for i in range(n):
  while stack and arr[stack[-1]] < arr[i]:
    answer[stack.pop()] = arr[i]
  stack.append(i)

for i in range(n):
  print(answer[i], end=' ')