오큰수 문제 보러가기

스택의 활용

• 단순히 원소값 자체를 저장해서 스택을 다루는게 아닌 원소의 인덱스값을 저장하므로써 스택을 이용할 수도 있음

문제

---

크기가 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)을 공백으로 구분해 출력한다.

---

핵심 포인트 요약

• A[i]의 오른쪽에 있는 수를 하나하나 비교해 A[i]보다 큰 수를 찾는 방법은 $O(N^2)$만큼의 시간복잡도가 발생해 시간초과를 초래함. 따라서 다른 방법을 생각해봐야함
  👉 원소의 인덱스를 저장하는 스택을 사용해 문제를 해결할 수 있음

---

코드

N = int(input())
A = list(map(int, input().split()))

result = [-1] * N	# -1로 초기화함(오큰수가 없는 경우 -1이기 때문)
stack = [0]			# A 리스트의 인덱스를 저장할 스택 초기화

for i in range(1, N):
    while stack and A[stack[-1]] < A[i]:
        result[stack.pop()] = A[i] 
    stack.append(i)

print(*result)

---

보충 설명

for i in range(1, N): 
    while stack and A[stack[-1]] < A[i]:
        result[stack.pop()] = A[i]
    stack.append(i)

👉 A = [3, 5, 2, 7] 일때 생각해보자

• 처음 스택에는 0이 들어가 있으며, i=1일때 A[stack[-1]]=3 와 A[1]=5의 원소를 비교한다. A[1]=5가 더크므로 result[0]에는 A[1]의 값인 5를 넣어준다. 즉, A에서의 0인덱스에 해당하는 첫번째 오큰수를 구한 것이다. 그리고 첫번째 오큰수를 구했기 때문에 pop하고 그 다음으로 구해야할 인덱스 1를 스택에 넣어준다.

---

혹시나 잘못된 부분이 있으면 댓글 부탁드립니다.