문제

문제

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

풀이

입력, 알고리즘, 출력 모두 신경써야 시간초과가 안 나오는 조금 까다로운 문제

입력

입력은 BufferedReader와 StringTokenizer로 받는다

출력

출력은 정답 배열을 원소마다 출력하면 시간이 오래 걸리기 때문에 StringBuilder에 받아서 최종 StringBuilder를 한번 출력해야 한다

알고리즘

알고리즘은 다음과 같다

입력값을 받은 입력배열, 정답을 저장할 정답배열, 그리고 알고리즘에 사용할 스택을 사용한다.

스택을 사용하는 이유는, 현재 인덱스보다 나중에 오는 인덱스의 현재 원소보다 큰 원소를 갖는, 가장 왼쪽 값을 찾아야 하기 때문이다. LIFO (Last-In, First-Out) 형태를 갖는 스택으로 관리하면, 아직 <오큰수>를 찾지 못한 인덱스는 스택에 계속 저장된다.

for loop을 iterate하면서 현재 원소가 스택의 가장 마지막 인덱스의 값(nums[stack.peek()])보다 크면, 정답 배열의 인덱스(ans[stack.peek()])에 현재 원소(nums[i])를 저장하고, stack에서 그 인덱스 값을 제거한다 (pop()).

코드

import java.util.*;
import java.io.*;
class Main {
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(st.nextToken());
        
        st = new StringTokenizer(br.readLine());
        
        int[] nums = new int[n];
        int[] ans = new int[n];
        for (int i = 0; i < n; i++) {
            nums[i] = Integer.parseInt(st.nextToken());
            ans[i] = -1;
            
        }
        
        Stack<Integer> stack = new Stack<>();
        
        for (int i = 0; i < nums.length; i++) {
            
            while (!stack.isEmpty() && nums[stack.peek()] < nums[i]) {
                int idx = stack.pop();
                ans[idx] = nums[i];
            }
            stack.push(i);
        }
        while(!stack.isEmpty()) {
    	    nums[stack.pop()] = -1;  
        }
        StringBuilder sb = new StringBuilder();
        for(int i = 0; i < n;i++){
            sb.append(ans[i]);
            sb.append(" ");
        }
        System.out.println(sb);
        
    }
}