[백준 17298번] 오큰수 (Python/파이썬)

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

17298번: 오큰수

특정 숫자보다 큰 수를 구하는 이번문제는 처음에 단순 while문으로 구현하였지만 역시나 시간초과가 떴다. 

다시 생각하다가 두 수의 크기를 비교하는 과정에서 큐/스택을 잘 활용하는 방안이 생각났다.

(백준 11003번 에서 큐/스택을 이용하여 최솟값을 찾아봤었다.)

오큰수가 없는 경우에는 -1을 출력해야하므로 답을 제출할 리스트에 -1을 미리 세팅해두었다.

 

예를 들어 오큰수를 구하고 싶은 리스트 a = [7, 4, 5, 10, 2, 4] 라고 하고, 임의의 스택을 deque라고 하자.
첫번째 수는 일단 스택에 추가한다. deque = [7]

 

두번째 수를 추가하기 전에 반복문으로 두번째수가 스택의 마지막 숫자보다 큰지 판단하여, 크다면 마지막 숫자의 오큰수는 추가할 숫자가 된다. 반복문이 끝나면 숫자를 스택에 추가한다.  7 > 4 이므로 스택에 추가해준다 deque = [7, 4]

 

세번째 수는 스택의 마지막 숫자인 4보다 세번째 숫자인 5는 4의 오큰수이다. 오큰수가 구해지면 스택에서 제거함과 동시에 해답 인덱스에 오큰수를 추가해준다.  7은 여전히 5보다 크므로 반복문이 종료되고 숫자를 추가한다. deque = [7, 5]

 

위의 과정을 계속 반복하면 오큰수가 모두 구해지지 않고 구해지지 않은 오큰수는 맨 처음에 세팅했던 -1이 출력된다.

코드는 다음과 같다.

 

import sys
from collections import deque
input = sys.stdin.readline
N = int(input())
A = list(map(int, input().split()))
answer = [-1 for _ in range(N)]
temp = deque([])


temp.append(0)
for i in range(1,N):
  if temp and A[temp[-1]] >= A[i]:
    temp.append(i)
    
  while temp and A[temp[-1]] < A[i]:
    answer[temp.pop()] = A[i]
  temp.append(i)

for i in answer:
  print(i, end = " ")