[백준] 17298.오큰수/Python - Gold4

❓문제

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

성능 요약

메모리: 205056 KB, 시간: 1140 ms

분류

자료 구조, 스택

 

문제 설명

크기가 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이다.

✍🏻풀이

오른쪽에 있는 수들 중 크면서 가장 왼쪽에 있는 수를 찾는 문제이다. 즉 현재 위치에서 바로 오른쪽에 있는 수가 오른쪽에 있는 수들 중 가장 왼쪽에 있는 수이다. 그러나 오른쪽에 있는 수가 항상 현재 수보다 큰 것은 아니다. 그러므로 바로 오른쪽 수가 작다면 스택에 현재 위치와 같이 넣는다. 다음 수로 넘어가 먼저 오른쪽에 있는 수가 현재 수보다 큰지 확인한다. 그리고 바로 오른쪽 수가 크다면 마찬가지로 이 수를 오큰수로 기록한다. 이 때, 스택에 아직 오큰수를 찾지 못한 수들이 있다면 현재 수의 오큰수를 스택에 있는 수들의 오큰수로 기록한다. 이 때, 스택에 있는 수가 현재 수의 오큰수보다 작아야 한다.

총 4개의 수 10, 8, 5, 9가 주어졌다고 가정해보자.
10의 오른쪽 수를 보면 8이므로 오큰수가 될 수 없다. 그러므로 스택에 (0, 10)을 넣는다. 넘어가서 8의 오른쪽 수는 7이다. 그러므로 8도 오큰수를 찾지 못했으므로 스택에 들어간다. 이제 5를 보면 5의 오른쪽 수는 9로 이 수는 오큰수가 된다. 이 때, 스택에 오큰수를 찾지 못하고 남아있는 수 8과 10이 있다. 9는 8의 오큰수가 된다. 그러나 9는 10의 오큰수가 될 수 없으므로 10은 오큰수를 찾지 못한다. 그러므로 10은 -1이 되고, 9는 오른쪽에 더 이상 수가 없으므로 오큰수를 찾을 수 없어 -1이 된다.
결과적으로 -1, 9, 9, -1이 된다.

💻코드

import sys

input = sys.stdin.readline

n = int(input())
A = list(map(int, input().split()))
st = []
res = [-1] * n
for i in range(n-1):
    if A[i] < A[i+1]:
        res[i] = A[i+1]
        while st and st[-1][1] < A[i+1]:
            res[st.pop()[0]] = A[i+1]
    else:
        st.append((i, A[i]))

print(*res)

📍후기

뭔가 아직 문제 풀 때 이 자료구조와 어떤 알고리즘을 써야 할지는 생각을 못 하고 난 그냥 최적만 생각하는 것 같다.