17298번: 오큰수

문제

크기가 N인 수열 A가 주어집니다. 수열의 각 원소 A_i에 대해서 오큰수 NGE(i)를 구하려고 합니다. A_i의 오큰수는 오른쪽에 있으면서 A_i보다 큰 수 중에서 가장 왼쪽에 있는 수를 의미합니다. 그러한 수가 없는 경우에 오큰수는 -1입니다.

• 입력: 첫째 줄에 수열 A의 크기 N (1 ≤ N ≤ 1,000,000)이 주어집니다. 둘째 줄에 수열 A의 원소들이 주어집니다.
• 출력: 총 N개의 수 NGE(1), NGE(2), ..., NGE(N)을 공백으로 구분해 출력합니다.

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

풀이: 스택을 활용한 O(N) 탐색

수열의 크기가 최대 1,000,000이므로 이중 반복문으로 모든 원소의 오른쪽을 확인하면 O(N^2)으로 시간 초과가 발생합니다. 스택(Stack)을 활용하여 시간 복잡도를 O(N)으로 단축해야 하는 전형적인 문제입니다.

1. 정답 배열 초기화: 오큰수를 찾지 못한 원소는 -1을 출력해야 하므로, 정답을 담을 ans 배열을 먼저 선언하고 Arrays.fill(ans, -1)을 이용해 모두 -1로 초기화해 둡니다.
2. 스택에 '인덱스' 저장: 스택에는 원소의 값이 아닌 '아직 오큰수를 찾지 못한 원소의 인덱스'를 저장합니다. 이를 통해 오큰수를 발견했을 때 정답 배열의 정확한 위치에 값을 기록할 수 있습니다.
3. 오큰수 탐색 로직: 수열을 0번 인덱스부터 순차적으로 탐색(for (int i = 0; i < n; i++))합니다.
   • 스택이 비어있지 않고, 스택의 맨 위 인덱스가 가리키는 값(a[stack.peek()])보다 현재 탐색 중인 값(a[i])이 크다면? -> 현재 값 a[i]가 스택 맨 위 원소의 오큰수가 됩니다.
   • 따라서 조건(a[stack.peek()] < a[i])을 만족하는 동안 계속 스택에서 인덱스를 꺼내고(pop()), 해당 인덱스의 정답 배열에 현재 값 a[i]를 기록해 줍니다.
   • 탐색 및 갱신이 끝나면 다음 오큰수를 찾아야 하므로 현재 인덱스 i를 스택에 넣습니다(push(i)).
4. 출력: 순회가 끝난 뒤 ans 배열을 그대로 출력하면 됩니다. 스택에 남아있는 인덱스들은 오큰수를 찾지 못한 원소들이며, 처음에 -1로 초기화해 두었기 때문에 별도의 후처리가 필요 없습니다.

코드

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

public 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());
       int[] a = new int[n];
       StringTokenizer st = new StringTokenizer(br.readLine());
       for (int i = 0; i < n; i++) {
          a[i] = Integer.parseInt(st.nextToken());
       }

       int[] ans = new int[n];
       // 오큰수가 없는 경우를 대비해 미리 -1로 가득 채움
       Arrays.fill(ans, -1);

       // 아직 오큰수를 찾지 못한 인덱스를 보관할 스택
       Stack<Integer> stack = new Stack<>();
       for (int i = 0; i < n; i++) {
          // 스택 맨 위 인덱스의 값보다 현재 값이 크다면 오큰수를 찾은 것
          while (!stack.isEmpty() && a[stack.peek()] < a[i]) {
             ans[stack.pop()] = a[i];
          }

          // 현재 인덱스를 스택에 푸시
          stack.push(i);
       }

       // 결과 출력
       for (int i : ans) {
          bw.write(i + " ");
       }

       bw.flush();
    }
}

알고리즘 분류

자료구조 (Data Structures), 스택 (Stack)