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

문제 (링크)

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

17298번: 오큰수

 

나의 풀이

import sys
from collections import deque

input = sys.stdin.readline
n = int(input())
arr = list(map(int, input().split()))
arr = deque(arr)

while arr:
    current = arr.popleft()
    exist = False

    for x in arr:
        if current < x:
            print(x, end=" ")
            exist = True
            break
    
    if not exist:
        print(-1, end=" ")

 

코드 설명 및 풀이법

스택으로 분류되어 있는 문제긴 했지만 큐를 사용해도 될 것 같았다.

 

매번 맨 왼쪽 값을 뽑아 숫자 리스트를 순서대로 돌면서 현재 popleft된 값보다 큰 값이 나오면 해당 값을 프린트하고 바로 break하도록 했다.

 

하지만 이 방식은 최악의 경우 O(n^2)만큼의 시간이 소요되었고, 시간초과 판정을 받았다.

 

이 방식은 어찌됐건 특정 요소보다 큰 값을 찾기 위해서 리스트를 모두 돌아야한다는 점에서 비효율적이다.

 

따라서, 스택을 이용해 아예 한 번의 반복문 안에 값들을 처리할 수 있도록 할 수 있다.

 

import sys
input = sys.stdin.readline

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

for i in range(n):
    try:
        while arr[stack[-1]] < arr[i]:
            result[stack.pop()] = arr[i]
    
    except:
        pass
    
    stack.append(i)

for i in range(n):
    print(result[i], end=" ")

 여기서 중요한 전제는, 만약 n번째 수가 n+1번째 수보다 커서 pop되지 못한 경우, n+2번째 수가 n번째 수보다 큰 상황이 오면 n+1번째 수 또한 n+2번째 수보다 작다는 것이다.

 

예를 들어서, 숫자 배열이 [5, 1, 7] 일 경우 첫 번째 숫자인 5는 두 번째 숫자 1보다 크므로 1은 오큰수가 될 수 없다. 

 

이 때, 스택에는 5와 1이 각각 담기게 되고, 다음 숫자인 7이 들어올 때 이 수는 첫 번째 숫자 5보다 크므로 5와 1 둘 다의 오큰수가 된다.

 

바로 이러한 성질을 스택을 통해서 구현한 것이다.