(백준) 17298.오큰수 / with.Stack

문제 링크: https://www.acmicpc.net/problem/17298

코드

import sys

n = int(sys.stdin.readline().rstrip())
a = list(map(int, sys.stdin.readline().rstrip().split()))

res = [-1] * n
stack = []

for i in range(n):  
    while stack and a[stack[-1]] < a[i]:
        res[stack.pop()] = a[i]
    stack.append(i)

print(*res)

풀이

1. 리스트 a가 [3, 5, 2, 7]로 주어질 때 각 인덱스의 오큰수는 [5, 7, 7, -1].
2. 결과를 보면 일단 마지막 인덱스의 오큰수는 무조건 -1 이라는 걸 알 수 있다.
   • 왜? 오큰수라는 것이, 자신보다 오른쪽에 있는 큰 숫자를 의미하기 때문이다. 그럼 결과로 출력할 list는 -1로 n만큼 초기화 하면 되겠다는 걸 알 수 있다.
3. 그리고 나머지 0부터 n-1까지 인덱스의 오큰수를 구할 땐, stack을 활용하면 된다.

• 근데 그냥 아래 코드처럼 이중 반복문으로도 구할 수 있지 않나? 왜 굳이 stack을 활용해야되지?

res = [-1] * n
for i in range (n):
  for j in range (i+1, n):
    if a[j] > a[i]:
      res[i] = a[j]
      break

• 물론 이렇게해도 답은 똑같이 나오겠지만, N의 최대 크기가 1,000,000이다. 위 코드처럼 이중반복문으로 풀면 시간복잡도가 O(logN^2)만큼 걸리기 때문에, 99.9%(?) 시간 초과가 날 것이 뻔하다.
• 따라서 stack을 활용해서 풀어야 한다. 그렇다면 stack을 활용해서는 어케 푸느냐?
• stack은 LIFO 후입선출, 즉 마지막에 들어온 아이템이 가장 먼저 나가는 특징을 갖고 있다. 마치 책을 column 방향으로 세워둔 다음 빼내는 것처럼 말이다.
• 따라서 stack을 일단 빈 리스트로 초기화 하고, n만큼 반복문을 돌려서 현재 숫자가 stack의 가장 상위에 있는 숫자보다 크다면, 그 숫자의 오큰수를 현재 숫자로 갱신한다.
• 이러면 뒤에서부터 하나씩 처리하며 현재 숫자가 스택에 쌓인 숫자들 보다 크다면 오큰수를 갱신할 수 있다.

디버깅

1. a = [3, 5, 2, 7]을 예로 들어보자.
2. 3 → 스택이 비어 있으므로 현재 인덱스 번호를 스택에 push. (stack = [0])
3. 5 → 3보다 크므로 3의 오큰수를 5로 설정. (stack = [], res = [5, -1, -1, -1])
   • 이후 5를 스택에 push. (stack = [1])
4. 2 → 5보다 작으므로 그냥 push. (stack = [1, 2])
5. 7 → 2보다 크므로 2의 오큰수를 7로 설정. (stack = [1], res = [5, -1, 7, -1])
   • 5보다 크므로 5의 오큰수를 7로 설정. (stack = [], res = [5, 7, 7, -1])
   • 이후 7을 스택에 push. (stack = [3])