백준 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이다.

---

코드

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

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());
        String s = br.readLine();
        StringTokenizer st = new StringTokenizer(s);
        Stack<Integer> stack = new Stack<>();
        int[] sequence=new int[n];

        for (int i = 0; i < n; i++)
            sequence[i]=Integer.parseInt(st.nextToken());

        int comparator=0;
        int count=0;
        for (int i = 1; i < n; i++) {
            if(sequence[comparator]<sequence[i]){
                sequence[comparator]=sequence[i];

                while(count>0&&sequence[stack.peek()]<sequence[i]){
                        sequence[stack.pop()]=sequence[i];
                        count--;
                }
            }
            else {
                stack.push(comparator);
                count++;
            }

            comparator++;
        }

        while(count>0){
            sequence[stack.pop()]=-1;
            count--;
        }
        sequence[sequence.length-1]=-1;

        for(int i:sequence)
            bw.write(i+" ");

        bw.flush();
    }
}

---

해결 과정

1. 단순히 수열의 0번 인덱스부터 본 뒤 그 수보다 큰 수를 선형적으로 탐색해서 찾으려면 최악의 경우 O(n^2)이 걸릴 것이고, 수열의 크기는 최대 100만이기 때문에 무조건 시간초과가 걸릴 것이다. 따라서 선형적으로 탐색하면서 조건에 맞는 수는 바로 출력해주고, 안 맞는 수는 Stack에 넣어준다.
   
   우선 다음 수를 봤을 때 이전 수보다 크면 나중에 답으로 출력해주기 위해 이전 수의 인덱스에 넣어준다. 그런데 만약 이전 수보다 작다면 스택에 넣어준다. 이렇게 해두면 바로 다음 수가 더 큰 수들은 이미 정답이 배열에 들어가있고, 다음 수가 더 작은 수들은 스택에 쌓여있을 것이다. 이 수들은 언제 탐색해주냐면, 다음 수가 더 큰 수가 나왔을 때 스택에 있는 수들 보다도 클 수 있기 때문에 스택의 top부터 비교해준다.

   왜 Queue가 아닌 Stack일까? 다음 수가 더 작다면 이전 수를 계속해서 패스한다. 이렇게되면 패스한 수들은 내림차순으로 된다. 그 어떤 경우도 오름차순으로 들어갈 수 없다. 따라서 제일 최근의 수와 비교해서 조건이 맞지 않으면 다른 패스한 수들도 비교할 필요도 없이 조건이 맞지 않다.

2. 시간초과!!(수열의 처음부터 끝까지 n회 반복되고, Stack과 비교할 때도 최소 1번, 최대 수열의 크기만큼 비교하지만 최악의 경우를 생각하더라도 처음부터 끝까지 패스하다가 마지막에 제일 큰 수가 나와서 n-1번 비교하게 된다. 최악의 경우에도 O(N) 인데 왜 시간초과인지 모르겠다)

3. 혹시 몰라서 BufferedWriter도 사용했고 스택을 검사할려고 할 때, 스택의 제일 위를 보기 위한 peek() 함수에서 스택이 Empty 상태라면 예외가 발생하기 때문에 isEmpty() 함수를 거쳤다. while문에서 조건 검사를 할 때마다 함수를 호출하기 때문에 시간이 걸리는 듯 하여 push할 때마다 count++ 해줬고, Stack.isEmpty()를 호출하지 않고 count가 양수면 검사하도록 바꿨다.

4. 😁