[백준] 17298번 - 오큰수 (Gold 4)

문제 링크

백준 17298번 - 오큰수 (Gold 4)

문제 설명

주어진 수열의 각 원소에 대한 오큰수를 출력한다.
오큰수란 해당 원소 $A_i$보다 오른쪽에 있으면서 $A_i$보다 큰 수 중 가장 왼쪽에 있는 수를 뜻한다.

주의점

• 시간복잡도에 유의해야한다.
  • $O(N^2)$으로는 시간 초과가 뜬다.

해결 전략

방법이 쉽게 떠오르지 않아 구글링을 통해 힌트를 얻고 풀었음을 밝힌다.

가장 쉽게 생각할 수 있는 방법은 각 원소마다 반복문을 돌려서 오큰수를 구하는 것이다.
하지만, 이렇게 접근하면 시간복잡도가 $O(N^2)$이 되어 시간 초과가 뜬다.

재귀적으로 오큰수를 구하는 과정을 생각해보자.
수식으로만 적으면 감이 잘 오지 않을테니, 수열 7, 3, 5, 8을 예시로 들어보자.

• $NGE(1)$
  • $A_2 \leq A_1$이므로, $NGE(2)$와 $A_1$을 비교한다.
  • NGE(2)
    • $A_3 > A_2$이므로, $NGD(2) = 5$
  • $NGE(2) \leq A_1$이므로, $NGE(3)$과 $A_1$을 비교한다.
  • NGE(3)
    • $A_4 > A_3$이므로, $NGD(3) = 8$
  • $NGE(3) > A_1$이므로, $NGE(1) = NGE(3) = 8$
• $NGE(4)$
  • 맨 끝 원소이므로 $NGE(4) = -1$

오큰수가 구해지는 순서를 담은 배열을 생각해보자.

• [1] : $NGE(1)$을 계산하고 싶다.
• [1 2] : $NGE(1)$을 계산하기 위해서 $NGE(2)$를 계산한다.
• [1] : $NGE(2)$가 계산되었다.
• [1 3] : $NGE(1)$을 계산하기 위해서 $NGE(3)$을 계산한다.
• [1] : $NGE(3)$이 계산되었다.
• [] : $NGE(1)$이 계산되었다.
• [4] : $NGE(4)$가 계산하고 싶다.
• [] : $NGE(1)$가 계산되었다.

구하고 싶은 오큰수의 인덱스를 순차적으로 담으면 맨 뒤에 해당하는 인덱스부터 오큰수가 계산됨을 볼 수 있다.
이는 구하고 싶은 오큰수의 인덱스를 스택으로 표현할 수 있음을 의미한다.

정답 코드

import java.io.*;
import java.util.*;

public class Main {

    static int a[];

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer st;
        StringBuilder sb = new StringBuilder();

        int N = Integer.parseInt(br.readLine());
        st = new StringTokenizer(br.readLine());
        a = new int[N];
        int NGE[] = new int[N];
        for (int i = 0; i < N; i++) {
            a[i] = Integer.parseInt(st.nextToken());
            NGE[i] = -1; // 뒤에 더 큰 수가 하나도 없는 경우의 기본값은 -1
        }
        Stack<Integer> indexStack = new Stack<>();
        for (int i = 0; i < N - 1; i++) {
            indexStack.add(i); // 인덱스를 스택에 추가
            // 구해야하는 인덱스의 오큰수가 a[i + 1]일 조건
            while (!indexStack.isEmpty() && a[i + 1] > a[indexStack.peek()]) {
                NGE[indexStack.pop()] = a[i + 1];
            }
        }
        for (int i = 0; i < N; i++) {
            sb.append(NGE[i]);
            sb.append(" ");
        }
        bw.write(sb.toString());
        bw.close();
    }
}

이 코드의 시간복잡도는 스택에 add하는 횟수와 pop하는 횟수를 보면 되는데, 각각이 $O(N)$이다.
따라서 전체 시간복잡도도 $O(N)$이 되어, $O(N^2)$에 비해 매우 빨라짐을 알 수 있다.