[백준 파이썬] 17298 오큰수

📌문제 출처

백준 단계별 문제풀이 - 스택 2
https://www.acmicpc.net/problem/17298

17298번: 오큰수

❓ 문제

❗ 풀이

• 스택과 인덱스 활용

📗 풀이 코드

import sys
input = sys.stdin.readline

n = int(input())
nums = list(map(int,input().split()))
ans = [-1]*n
stack = [] # 오큰수를 찾아야하는 nums의 인덱스를 담아준다.

for i in range(n):
    # stack이 있고, 오큰수를 찾아야하는 수보다 더 큰 수가 나타났다. => 오큰수를 발견했다.
    while stack and nums[stack[-1]] < nums[i] :
        ans[stack.pop()] = nums[i] # 오큰수를 발견한 인덱스는 pop으로 빼내준다.
    stack.append(i) # 이때 i는 오큰수가 되는 인덱스기에 이 인덱스에 대한 오큰수도 찾아줘야한다.
print(*ans)

📗 코드 해설

• nums의 모든 수마다 반복문을 돌리며 오큰수를 찾아주면 O(n**2)의 시간복잡도가 나와 실패한다.
• nums의 인덱스를 돌며 오큰수를 찾을 때마다 다시 이전의 인덱스를 하나씩 살펴보며 찾은 오큰수가 다른 수의 오큰수가 될 수 있을지를 확인한다.
• ans에는 오큰수가 없을 때 정답을 출력하도록 -1을 n개 만큼 넣어둔다.
  stack에는 오큰수를 찾아야하는 nums의 인덱스를 담아줄 것이다.
• 먼저 0~n-1까지 for 반복문으로 돌며 nums의 수를 꺼낸다. 이때 nums[i]는 오큰수가 될 수 있는 수라고 생각하자.
  이때 stack이 비어있다면 오큰수를 찾아야하는 nums의 인덱스가 없는 것이므로 i를 담아준다.
  이 경우 nums[i]가 오큰수가 될 수 없는 경우를 아직 못 찾은 경우이고,
  stack에 들어간 이후로 i는 오큰수를 찾아야하는 인덱스가 된다.
• 만일 stack이 있고(오큰수를 찾아야하는 인덱스가 있고), nums[i]가 해당 인덱스의 오큰수가 될 수 있다면
  오큰수를 발견한 인덱스는 다시 한번 오큰수를 찾아줄 필요가 없으므로 stack.pop()으로 빼내주고,
  해당 인덱스의 정답 오큰수를 ans에 넣어준다.
• 그리고 다시 while문을 돌면서 nums[i]가 다른 인덱스(앞선 인덱스)들의 오큰수가 될 수 있는지 확인한다.
  이를 통해 하나의 값이 다른 앞선 인덱스들의 오큰수가 될 수 있는 경우 탐색 속도를 높여준다.
• while문이 끝나면 stack이 비어있는 경우와 마찬가지로 더이상 nums[i]가 오큰수가 될 수 없는 것이므로
  i의 상황을 바꾸어 오큰수를 찾아야하는 인덱스로 바꾸어준다.