[JAVA] 백준 17298번: 오큰수 (초급 1-9)

문제

크기가 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)을 공백으로 구분해 출력한다.

 

문제 풀이

일단 이 문제도 Scanner를 이용해서 문제를 풀면 시간초과가 발생한다. BufferReader + StringBuilder를 사용하자.

for문 2개를 사용해서 문제를 풀면되겠다 생각했는데, 바로 시간초과가 발생.

for문 1개로 변환하는 방법을 생각해보았다.

stack을 이용하는 것이다.

 

1. for문으로 숫자 1개씩 확인한다.

2. stack이 비어있을 경우 valueStack에 현재 값을 넣어주고, indexStack에 현재 값의 index를 넣어준다.

3. stack이 비어있지 않을 경우, stack의 마지막 값이 for문으로 입력 받은 수 보다 클 경우 결과 값으로 for문으로 입력받은 값을 저장한다. (while문) -> valueStack에 현재 값을 넣어주고, indexStack에 현재 값의 index를 넣어준다.

위 로직을 코드로 확인해보자.

for (int i = 0; i < n; i++) {
    while (!valueStack.isEmpty() && valueStack.peek() < inputArray[i]) {
        resultArray[indexStack.pop()] = inputArray[i];
        valueStack.pop();
    }
    valueStack.push(inputArray[i]);
    indexStack.push(i);
}

 

1. 위와 같은 로직이 가능한 이유는 valueStack에는 점점 작은수가 들어오기 때문이다. (중요!!!!)

1-1) valueStack = {4, 2, 1} (가능)

1-2) valueStack = {4, 5} (불가능)

왜냐?? 우리가 valueStack 마지막에 들어가있는 값 보다 작은 수만 valueStack에 넣기 때문이다.

이 사실을 이용하여 코드를 작성한 것이다.

 

2. 사실 valueStack은 필요가 없다. valueStack.peek() 부분을 inputArray[indexStack.peek()] 으로 변경하면 된다.

 

전체 코드를 확인해보자.

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Stack;

public class baekjoon17298_3 {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringBuilder sb = new StringBuilder();

        Stack<Integer> valueStack = new Stack<>();
        Stack<Integer> indexStack = new Stack<>();

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

        String[] temp = br.readLine().split(" ");
        for (int i = 0; i < n; i++) {
            inputArray[i] = Integer.parseInt(temp[i]);
        }

        for (int i = 0; i < n; i++) {
            while (!valueStack.isEmpty() && valueStack.peek() < inputArray[i]) {
                resultArray[indexStack.pop()] = inputArray[i];
                valueStack.pop();
            }
            valueStack.push(inputArray[i]);
            indexStack.push(i);
        }

        while (!valueStack.isEmpty()) {
            resultArray[indexStack.pop()] = -1;
            valueStack.pop();
        }

        for (int i : resultArray) {
            sb.append(i + " ");
        }
        System.out.println(sb);
    }
}