(17298) 오큰수 (Java)

문제

크기가 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

출처

문제를 만든 사람: baekjoon
데이터를 추가한 사람: rhdqor213

알고리즘 분류

자료 구조
스택

Solution

1 시간초과

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

public class Main {
    public static int n;
    public 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));

        n = Integer.parseInt(br.readLine());
        a = new int[n];

        StringTokenizer st = new StringTokenizer(br.readLine());

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

        for(int i = 0; i < n; i++) {
            bw.write(String.valueOf(findNge(a[i], i)) + " ");
        }
        bw.flush();
        bw.close();
        br.close();
    }

    public static int findNge(int x, int idx) {
        if(idx == n) return -1;

        for(int i = idx+1; i < n; i++) {
            if(x < a[i]) {
                return a[i];
            }
        }

        return -1;
    }
}

2

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

public class Main {
    public static int n;
    public static int[] a;
    public static int[] nge;
    public static Stack<Integer> stk;

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

        n = Integer.parseInt(br.readLine());
        a = new int[n];
        nge = new int[n];
        stk = new Stack<Integer>();

        StringTokenizer st = new StringTokenizer(br.readLine());

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

        // 오큰수 찾기
        findNge();

        // 정답 출력
        for(int i = 0; i < nge.length; i++) {
            bw.write(String.valueOf(nge[i]) + " ");
        }

        bw.flush();
        bw.close();
        br.close();
    }

    public static void findNge() {
        // n개의 수를 돌면서
        for(int i = 0; i < n; i++) {
            // 스택에 a배열의 인덱스가 들어있는 경우
            while(!stk.isEmpty()) {
                // 현재 값과 스택에 들어있는 인덱스를 갖는 배열 값을 비교하여
                // 배열의 값이 큰 경우 오큰수이므로, nge[스택pop인덱스]에 해당 값을 할당한다.
                if(a[stk.peek()] < a[i]) {
                    nge[stk.pop()] = a[i];
                // 현재 값보다 스택에 들어있는 인덱스를 갖는 배열의 값이 더 크거나 같은 경우
                // 오큰수가 아니므로 현재 값을 스택에 담아주고 다음 인덱스를 탐색한다.
                } else {
                    break;
                }
            }
            stk.push(i);
        }

        // 수열을 모두 탐색하였는데도 스택에 남아있는 수의 경우
        // 오큰수가 존재하지 않는다는 것이므로 -1을 할당한다.
        while(!stk.isEmpty()) {
            nge[stk.pop()] = -1;
        }
    }
}

Feedback

Solution 1. 시간초과가 발생하였다. 최악의 경우 수열의 크기가 1,000,000만이고 A1 = 1,000,000, A 100만번째 수가 1인 경우, 이중반복문의 시간복잡도 O(N^2)은 100만번, 999,999번, ..., 1번으로 총 500,000,500,000회 수행하게 된다.

Solution 2. 스택을 사용하여 해결하였다. 수열(a)의 첫번째 수부터 스택에 쌓고, 수열을 돌면서 스택의 탑과 수열의 값을 비교한다. 이때 수열의 값이 스택의 탑보다 큰 경우 오큰수이므로 nge 배열에 해당 수열의 값을 저장한다. (스택에는 값 자체를 담고있지 않고, 수열의 인덱스를 담고있기 때문에 가능)

결과적으로 스택에 쌓여있는 수들은 아직 오큰수를 발견하지 못한 수들만 존재할 것이고, 수열의 마지막 수까지 탐색이 완료되었을 때 스택에 남아있는 수는 끝까지 오큰수를 발견하지 못한 수이므로 mge 배열에 -1을 담아준다.