오큰수

문제

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

풀이

각 위치에서 오른쪽을 끝까지 탐색하면 최대 백만 원소에 대해 O(N²)이 된다. 새 숫자를 왼쪽부터 읽을 때, 아직 오큰수를 못 찾은 인덱스만 스택에 둔다. 현재 수 A[i]가 스택 꼭대기의 수보다 크면, A[i]가 그 인덱스에서 처음 만난 더 큰 오른쪽 수다. 꼭대기를 꺼내 답을 채우고, 더 이상 꺼낼 수 없을 때 현재 인덱스를 넣는다.

[3,5,2,7]을 보면 3의 인덱스가 먼저 쌓인다. 5를 읽는 순간 3을 꺼내 답 5를 적는다. 2는 5보다 작아 쌓인다. 마지막 7은 2와 5를 차례로 해결하고, 자기 자신은 해결되지 않은 채 남아 -1이 된다. 스택에는 아래에서 위로 값이 내림차순 또는 같은 값으로 남는다. 같은 값은 ‘큰’ 수가 아니므로 비교는 <여야 한다.

각 인덱스는 한 번 들어가고 최대 한 번 나오므로 시간 O(N), 수열·답·스택 공간 O(N)이다. N이 백만이어서 split(" ")로 백만 개 문자열을 만들면 메모리 부담이 크다. 아래 코드는 바이트 입력과 정수 배열 스택을 사용한다. 모두 같은 수인 경우 답은 전부 -1이며, 엄격한 비교를 확인하는 경계 사례다.

소스코드

import java.io.BufferedInputStream;
import java.io.BufferedWriter;
import java.io.OutputStreamWriter;
import java.util.Arrays;

public class Main {
    public static void main(String[] args) throws Exception {
        FastScanner in = new FastScanner();
        int n = in.nextInt();
        int[] numbers = new int[n];
        int[] answer = new int[n];
        int[] stack = new int[n]; // 아직 오큰수를 못 찾은 인덱스
        Arrays.fill(answer, -1);
        int size = 0;

        for (int i = 0; i < n; i++) {
            numbers[i] = in.nextInt();
            while (size > 0 && numbers[stack[size - 1]] < numbers[i]) {
                answer[stack[--size]] = numbers[i];
            }
            stack[size++] = i;
        }

        StringBuilder result = new StringBuilder(n * 3);
        for (int i = 0; i < n; i++) {
            if (i > 0) result.append(' ');
            result.append(answer[i]);
        }
        BufferedWriter out = new BufferedWriter(new OutputStreamWriter(System.out));
        out.write(result.toString());
        out.newLine();
        out.flush();
    }

    static class FastScanner {
        private final BufferedInputStream in = new BufferedInputStream(System.in);
        private final byte[] buffer = new byte[1 << 16];
        private int length = 0, position = 0;

        int read() throws Exception {
            if (position >= length) {
                length = in.read(buffer);
                position = 0;
                if (length == -1) return -1;
            }
            return buffer[position++];
        }

        int nextInt() throws Exception {
            int ch;
            do { ch = read(); } while (ch <= ' ' && ch != -1);
            int value = 0;
            while (ch > ' ') {
                value = value * 10 + ch - '0';
                ch = read();
            }
            return value;
        }
    }
}