[JAVA] 백준 (골드4) 17298번 오큰수

링크

https://www.acmicpc.net/problem/17298

---

문제 설명

정답률 34.206%
크기가 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)이 주어진다.

4
3 5 2 7

---

출력 예제

• 총 N개의 수 NGE(1), NGE(2), ..., NGE(N)을 공백으로 구분해 출력한다.

5 7 7 -1

---

풀이

N의 최대값이 $10^6$으로 반복문으로 오큰수를 찾을 경우 시간 초과가 발생하므로 스택을 이용하여 풀이한다.

스택에 새로 들어오는 수가 스택의 top보다 크면 그 수는 top의 오큰수가 되며, 오큰수가 존재하지 않으면 모두 -1로 할당한다. 여기서 중요한 점은 스택에 수열의 값이 아닌 인덱스를 저장한다.

스택에 먼저 0번 인덱스를 추가해두고 1부터 N-1까지 반복한다. 예제를 기준으로 생각해보면 수열은 seq = {3, 5, 2, 7} 이다. 반복 과정은 다음과 같다.

• i = 1
  • stack = [0]
  • seq[top] < seq[i]
  • NGE[top] = seq[i]
  • pop()
    • stack = []
  • push(i)
    • stack = [1]
• i = 2
  • stack = [1]
  • seq[top] > seq[i]
  • push(i)
    • stack = [1, 2]
• i = 3
  • stack = [1, 2]
  • seq[top] < seq[i]
  • NGE[top] = seq[i]
  • pop
    • stack = [1]
  • seq[top] < seq[i]
  • NGE[top] = seq[i]
  • pop
    • stack = []
  • push(i)
    • stack = [3]

오큰수를 구한 뒤 스택에 원소가 남아 있을 경우 이 원소에 해당하는 수열의 오큰수는 -1로 할당한다.

• stack = [3]
• NGE[3] = -1
  

코드

//백준
public class Main {

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

        System.setIn(new FileInputStream("src/input.txt"));
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        int N = Integer.parseInt(br.readLine());  //10^6
        int[] seq = new int[N];  //수열
        StringTokenizer st = new StringTokenizer(br.readLine());
        for (int i = 0; i < N; i++) {
            seq[i] = Integer.parseInt(st.nextToken());  //10^6
        }
        int[] NGE = new int[N];  //오큰수

        Stack<Integer> stack = new Stack<>();
        stack.push(0);

        for (int i = 1; i < N; i++) {

            //스택의 top보다 수열의 원소가 크고 스택이 빌 때까지 반복
            while (!stack.isEmpty() && seq[stack.peek()] < seq[i]) {
                //스택의 top의 오큰수를 수열의 원소로 할당
                NGE[stack.pop()] = seq[i];
            }
            stack.push(i);
        }

        //오큰수가 없는 원소들은 -1
        while (!stack.isEmpty()) {
            NGE[stack.pop()] = -1;
        }

        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < N; i++) {
            sb.append(NGE[i]).append(" ");
        }

        System.out.println(sb);
    }
}