N의 최대 크기가 1,000,000이므로 반복문으로 오큰수를 찾으면 제한 시간을 초과하게 된다. 따라서 이 문제는 스택을 이용하여 해결해야한다.
예제의 [3, 5, 2, 7]에 대한 오큰수를 구하려면 3의 오큰수인 5, 5의 오큰수인 7, 2의 오큰수인 7, 7은 오큰수가 없어 -1이다.
핵심은 스택에 탐색중인 인덱스를 저장한 후, 탐색중인 입력값과 비교한다. 오큰수가 맞을 경우 스택에 저장된 인덱스 값을 정답배열의 인덱스로 활용하여 오큰수의 정답을 저장하는 원리이다.
문제 풀이 방법은 다음과 같다.
-
입력값이 스택의 top에 위치한 인덱스에 해당하는 입력값보다 큰 경우 (== 오큰수인 경우)
e.g. [3, 5, 2, 7]에서 입력값 A[1] = 5, 스택의 top에 위치한 인덱스에 해당하는 입력값 = 3=> 스택에서 pop을 통해 인덱스를 뺀 다음 answer의 인덱스로 사용하여 오큰수 저장
e.g. 스택 top에 위치한 인덱스에 해당하는 입력값(3)의 인덱스 = 0을 사용한다.
answer[0] = A[1] # 오큰수 저장=> 입력값 인덱스인 i를 스택에 추가
e.g. 예제를 그대로 이어가면 push(1)이 됨
-
입력값이 스택의 top에 위치한 인덱스에 해당하는 입력값보다 작거나 같은 경우 (== 오큰수가 아닌 경우)
=> 입력값 인덱스인 i를 스택에 추가
위 1, 2번 과정을 N번 반복
- 스택에 남아있는 값을 처리
이 경우는 배열의 오른쪽에 자신보다 큰 수가 없는 경우이므로 -1 저장
Python
import sys
input = sys.stdin.readline
N = int(input())
A = list(map(int, input().split()))
answer =[0] * N # 결과를 저장할 배열
stack = [] # 현재 탐색 중인 인덱스를 저장
for i in range(N):
# 스택이 비어있지 않고, 입력값이 스택의 top에 위치한 인덱스에 해당하는 입력값보다 큰 경우
# top에 위치한 스택을 pop하여 answer에 현재 값을 저장한다
# 정리하면 오큰수를 발견한 경우 while문을 통해 오큰수의 조건을 만족할때까지 반복하면서
while stack and A[stack[-1]] < A[i]:
# 스택에서 인덱스를 pop한 것을 answer 인덱스 위치에 오큰수의 정답으로 저장
answer[stack.pop()] = A[i]
# 입력값 인덱스인 i를 스택에 추가
stack.append(i)
# 스택에 남은 값 처리
while stack:
# 이 경우는 오른쪽에 자신보다 큰 숫자가 없는 경우
answer[stack.pop()] = -1
for i in range(N):
print(str(answer[i]), end=" ")Java
import java.io.*;
import java.util.Stack;
import java.util.StringTokenizer;
public class P_17298 {
public static void main(String[] args) throws IOException {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
int N = Integer.parseInt(br.readLine());
int[] A = new int[N];
int[] Result = new int[N];
StringTokenizer st = new StringTokenizer(br.readLine());
for (int i = 0; i < N; i++) {
A[i] = Integer.parseInt(st.nextToken());
}
Stack<Integer> stack = new Stack<>();
stack.push(0);
for (int i = 1; i < N; i++) {
while (!(stack.isEmpty()) && A[i] > A[stack.peek()]) {
Result[stack.pop()] = A[i];
}
stack.push(i);
}
while (!(stack.isEmpty())) {
Result[stack.pop()] = -1;
}
BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
for (int i = 0; i < N; i++) {
bw.write(Result[i] + " ");
}
bw.write("\n");
bw.flush();
}
}※ 오큰수의 조건검사가 while인 이유?
예제인 [3, 5, 2, 7]에서 3의 오큰수인 5는 구했고, 5의 오큰수를 구해야 한다고 가정해보자.
2는 오큰수가 아니므로 stack.append(i)가 될 것이고, 그 이후에 7을 만나게 되며 오큰수를 발견한다.
현재 스택의 top에 있는 숫자는 가장 최근에 append한 2이고 이를 pop하여 오큰수를 7이라고 저장한다.
그리고 반복하여 2 밑에 깔려있는 5 역시 7을 오큰수로 저장한다.
if 문을 사용하면 가장 최근에 append한 2에 맞는 오큰수만 찾고, 이후 스택 밑에 깔린 수는 스킵하게 되므로 while문을 사용해야 한다.
