[Baekjoon/Programmers] 17298 오큰수

17298 오큰수문제 해결 과정.

수열이 주어졌을 때, 한 원소에서 오른쪽에 위치한 원소들 중에 원소보다 큰 수 중 가장 작은 수를 해당 원소의 오큰수라고 한다. 예를 들어 3 5 7 2 의 수열이 주어졌을 때, 3의 오큰수는 5, 5의 오큰수는 7이다. 단순히 한 원소의 오른쪽을 탐색하고 원소보다 큰 수들 중에서 가장 작은 수를 탐색하는 문제같지만 이런 방식으로 풀면 시간복잡도가 매우 커진다.

그렇다면 뒤에서부터 오큰수를 찾아나가면 어떨까?

import java.io.*;
import java.util.stream.Stream;

public class Main {

    public static void main(String[] args)  throws IOException {

        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

        int n = Integer.parseInt(br.readLine());
        int[] sequence = Stream.of(br.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();
        StringBuilder ngeReverse = new StringBuilder();

        int[] nge = new int[n]; nge[n-1] = -1;
        int maxNge = -1;

        for (int i = n-1; i>0; i--) {
            if (sequence[i-1]>sequence[i]) {
                if (sequence[i-1]>nge[i]) {
                    if(sequence[i-1]>maxNge) nge[i-1]=-1;
                    else nge[i-1]=maxNge;
                }
                else nge[i-1] = nge[i];
            } else if (sequence[i-1] < sequence[i]) {
                if (sequence[i-1]<nge[i]) nge[i-1] = Math.min(nge[i],sequence[i]);
                else nge[i - 1] = sequence[i];
            } else nge[i-1] = nge[i];

            if (maxNge<nge[i]) maxNge = nge[i];
        }

        for (int i=0; i<n; i++) {
            bw.write(nge[i]+"");
            if (i<n-1) bw.write(" ");
        }
        bw.flush();
        bw.close();
    }
}

수열의 마지막 수의 오큰수를 -1로 지정한 뒤, 뒤에서부터 인접한 값과 비교하는 방식으로 코드를 작성했다. i-1번째와 i번째의 값을 비교해 앞에 있는 값이 더 큰 경우 오큰수보다 크면 -1, 작으면 해당값의 오큰수를 넣어주었고. 앞에 있는 값이 더 작은 경우 오큰수보다도 작으면 오큰수와 해당값 중 더 작은 값을 오큰수로 하고 오큰수보다는 크다면 값을 입력하는 식으로 작성했다. 그러나 이는 같은 경우를 고려하지 않은 설계이고 이웃하는 값만 고려해 오큰수가 있어도 -1을 출력하는 경우가 있었다. 오큰수의 max값을 저장하는 변수를 추가해 해당 반례를 해결했으나 3 1 2 4 5 라는 수열이 들어왔을 때, 이 Max값 변수로 인해 5 2 4 5 -1 라는 결과가 출력됐다. 오큰수 내에서 해결하는 건 결국 또 반복문을 중첩해 사용해야 하기 때문에 이러한 로직으로 문제를 해결하는 것에는 무리가 있을 거라는 결론을 내렸다.

결국 스택을 사용해서 해결해야 한다. 처음에는 수열에 있는 값을 차례로 push하면서 기존에 있던 값과 push하려는 값을 비교해 push하는 값이 크다면 기존 값을 pop하는 방식으로 구현하려 했으나 이렇게 하면 stack에 남아 있는 값들의 index를 알아 오는 과정이 꽤 까다로울 것 같았다.

이는 인덱스를 스택에 넣으면 해결되었다!

오큰수를 찾지 못한 값들의 인덱스를 스택에 넣어두고, 더 큰값이 들어오면 pop을 수행한다. 끝까지 stack에 남아 있는 값들은 오큰수가 없는 값들이므로 오큰수에 -1을 넣어주는 방식으로 코드를 작성했더니 해결되었다.

import java.io.*;
import java.util.Stack;
import java.util.StringTokenizer;

public class Main {

    public static void main(String[] args) throws IOException {

        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

        int N = Integer.parseInt(br.readLine());
        int[] sequence = new int[N];
        int[] nge = new int[N];

        StringTokenizer st = new StringTokenizer(br.readLine());
        for (int i = 0; i < N; i++) {
            sequence[i] = Integer.parseInt(st.nextToken());
        }

        Stack<Integer> index = new Stack<>();
        index.push(0);
        
        for (int i = 1; i < N; i++) {
            while(!index.empty() && (sequence[i] > sequence[index.peek()])) {
                nge[index.pop()] = sequence[i];
            }
            index.push(i);
        }

        while(!index.empty()) {
            nge[index.pop()] = -1;
        }

        StringBuilder answer = new StringBuilder();
        for(int n: nge) {
            answer.append(n).append(" ");
        }
        System.out.println(answer.toString().trim());
    }
}

1등으로 푼 사람이 380ms여서 자괴감이 살짝 들었지만 첫 골드 문제를 해결하는 것에 의의를 두기로 했다… ^___^ 오등큰수는 더 빨리 풀어봐야지….