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

백준 17298 - 오큰수

시간 제한 1초, 메모리 제한 512MB

# 조건

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

입력

• 첫째 줄에 수열 A의 크기 N (1 ≤ N ≤ 1,000,000)이 주어진다. 둘째 줄에 수열 A의 원소 A1, A2, ..., AN (1 ≤ Ai ≤ 1,000,000)이 주어진다.

출력

• 총 N개의 수 NGE(1), NGE(2), ..., NGE(N)을 공백으로 구분해 출력한다.

  # 접근 방법

• 오큰수란 현재 수에서 오른쪽으로 순회하며 가장 먼저 나오는 나보다 큰 수로 이해하면 된다.

• 따라서 스택을 이용해서 풀어준다.

• for문으로 nums리스트를 순회하며 stack이 비어있지 않고, stack의 top이 현재 수보다 작은 경우 while문을 실행해준다.

  • 즉, 현재 기준으로 삼은 숫자보다 큰 숫자가 왼쪽에 나온다면 그 이전의 숫자는 자신들의 오큰수를 찾은 것이다.
  • 또한, stack이 비어있더라도 위의 로직과 같이 이미 오큰수를 찾아서 stack에서 pop되었기 때문이다!

import sys  
sys.stdin = open('input.txt')  
input = sys.stdin.readline  


N = int(input())  
nums = [*map(int, input().split())]  
ans = [-1] * N  
s = []  
for i in range(N):  
    while s and nums[s[-1]] < nums[i]:  
        ans[s.pop()] = nums[i]  
    s.append(i)  

print(*ans)