문제
https://www.acmicpc.net/problem/17298
크기가 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)을 공백으로 구분해 출력한다.
예제 입력 1
4
3 5 2 7
예제 출력 1
5 7 7 -1
예제 입력 2
4
9 5 4 8
예제 출력 2
-1 8 8 -1
접근
이런 유형의 모노토닉 스택 문제는 이미 몇 번 접해본 적이 있어서 스택을 써야 한다는 생각은 어렵지 않게 할 수 있었다.
수열의 원소를 하나씩 순회하면서 아래 단계를 반복해 수행한다.
1. 현재 원소가 스택 맨 위에 있는 수보다 크다면 현재 원소가 스택 맨 위에 있는 수의 오큰수이다. 오큰수를 찾았으므로 스택 맨 위에 있는 수를 pop한다.
2. 스택이 비게 되거나 현재 원소가 스택 맨 위에 있는 수보다 크지 않을 때까지 1번을 반복하며 현재 원소가 오큰수인 수들을 스택에서 찾는다.
3. 현재 원소가 오큰수인 수들을 모두 찾았다면 스택에 현재 원소를 push한다.
수열의 원소를 모두 순회했음에도 pop되지 않고 남아 있는 숫자들이 있다면, 해당 숫자들은 수열 내에서 오큰수가 발견되지 않았다는 뜻이므로 -1이 오큰수가 될 것이다.
다른 문제와의 차이점이라면 오큰수를 순서대로 하나씩 출력해 줘야 한다는 거였다. 스택에서 push와 pop을 반복하다보면 오큰수가 처음 수열의 순서대로 발견되지 않을 가능성이 매우 높은데, 이 순서를 어떻게 관리할 것인가를 좀 고민했다.
구현
처음에는 Element라는 객체에 원소의 값(value), 원소의 인덱스 값(index), 원소의 오큰수 정보(rightMax)를 저장한 후에 원소의 인덱스 값 순으로 정렬해 출력하는 방식으로 구현했다.
import java.io.*;
import java.util.*;
import java.util.stream.Collectors;
class Main {
public static void main(String[] args) throws IOException {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
int N = Integer.parseInt(br.readLine());
StringTokenizer st = new StringTokenizer(br.readLine());
Stack<Element> q = new Stack<>();
List<Element> popedElements = new ArrayList<>();
for (int i = 0; i < N; i++) {
int currVal = Integer.parseInt(st.nextToken());
while (!q.isEmpty() && q.peek().value < currVal) {
Element poped = q.pop();
poped.rightMax = currVal;
popedElements.add(poped);
}
q.push(new Element(currVal, i + 1));
}
while (!q.isEmpty()) {
popedElements.add(q.pop());
}
popedElements.sort((o1, o2) -> o1.index - o2.index);
System.out.println(String.join(" ", popedElements.stream().map(Element::getRightMax).collect(Collectors.toList())));
}
static class Element {
final int value;
final int index;
int rightMax;
Element(int value, int index) {
this.value = value;
this.index = index;
rightMax = -1;
}
String getRightMax() {
return String.valueOf(rightMax);
}
}
}하지만 너무 느린 것 같아서... 풀이를 찾아보니 스택에 원소 값을 직접 push하는 대신 인덱스 값을 push하고 값 비교를 할 때는 원 배열인 A에서 해당 인덱스 값을 참조하는 식으로 구현하면 인덱스 값을 저장하기 위해 따로 메모리를 추가로 할당할 필요가 없더라...!
그리고 A의 오큰수를 찾은 후에는 A[i]의 값은 더 이상 사용되지 않기 때문에 오큰수 값을 저장하는 배열이나 리스트를 따로 생성하는 대신 A[i]에 A[i]의 오큰수 값을 덮어 씌워 저장하면 메모리 사용을 더 줄일 수 있었다. 그리고 이렇게 하면 배열 A에 오큰수가 자연스럽게 차례대로 저장되기 때문에 정렬을 할 필요가 없어 시간 복잡도도 감소한다.
import java.io.*;
import java.util.*;
class Main {
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];
StringTokenizer st = new StringTokenizer(br.readLine());
Stack<Integer> q = new Stack<>();
for (int i = 0; i < N; i++) {
A[i] = Integer.parseInt(st.nextToken());
while (!q.isEmpty() && A[q.peek()] < A[i]) {
A[q.pop()] = A[i];
}
q.push(i);
}
while (!q.isEmpty()) {
A[q.pop()] = -1;
}
StringBuilder sb = new StringBuilder();
for (int e : A) {
sb.append(e).append(" ");
}
System.out.println(sb.toString());
}
}
아래가 내가 처음 구현한 풀이, 위가 최적화한 풀이의 채점 결과이다.
