문제 출처

오큰수 풀이

스택으로 풀 수 있는 문제의 한 종류인 것 같아서 따로 정리한다.

직관

1. answer 배열을 [0, 0, 0, ...]으로, stack은 빈 리스트로 initialize한다.
2. 문제의 배열(arr이라 하자)을 순서대로 순회하면서 "stack의 head에 있는 값"보다 "arr에서 보고 있는 값"이 더 작거나 같으면 stack에 push한다.
   • 이때 stack에는 배열의 값이 아니라 index를 넣어야 한다. 이유는 4번에서 설명한다.
3. stack의 head에 해당하는 값이 더 크다면 그 값을 pop한다. 이후 그 값에 해당하는 위치에 answer를 업데이트한다.
4. 배열을 다 순회안 이후 stack에 남아있는 index에 해당하는 answer배열의 값들을 $-1$로 처리한다. 이 처리를 위해 stack에 index를 넣었던 것이다!

글로 쓰니 더 헷갈리는 듯....코드는 의외로 간단하다.

코드

import sys


def main():
    input = sys.stdin.readline
    N = int(input())
    arr = list(map(int, input().split()))
    answer = [0] * N
    stack = []
    for i in range(N):
        num = arr[i]
        while stack and arr[stack[-1]] < num:
            answer[stack.pop()] = num
        stack.append(i)
    for j in stack:
        answer[j] = -1
    print(*answer)


if __name__ == "__main__":
    main()

핵심

이 문제의 핵심은 stack 배열의 특징이다. 어떠한 경우에도 stack은 내림차순의 순서를 지키게 된다. 따라서 하나의 while문 안에서 순서대로 pop하면서 보는 것이 가능하다. 지금 pop하는 값은 이전에 pop하는 값보다 작거나 같았을 것이므로!!