백준 17298 - 오큰수

YongHyun·2023년 3월 31일

백준, 프로그래머스 - 코딩테스트

목록 보기

12/43

문제

시간제한 1초

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

처음 문제를 풀때는 원소의 값 하나를 지정하고(변수명 find) 그 값을 제외한 오른쪽에 나머지 모든 수를 스택에 넣어준 후 find보다 큰 값이면 출력해주는 형식으로 접근했다.

하지만 문제에서는 시간제한은 1초이고 N의 최대 크기가 1,000,000이므로 위와 같은 방식으로 풀면 시간초과가 나온다.

시간초과가 나온 잘못된 풀이

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

class Main {

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

        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        Stack<Integer> stack = new Stack<>();
        StringBuilder sb = new StringBuilder();
        int N = Integer.parseInt(br.readLine());
        int A[] = new int[N];
        boolean findMax = false;

        StringTokenizer st = new StringTokenizer(br.readLine());
        for(int i = 0; i < N; i++)
            A[i] = Integer.parseInt(st.nextToken());

        for(int i = 0; i < N - 1; i++) {
            int find = A[i];
            for(int j = N - 1; j > i; j--){
                stack.push(A[j]);
            }

            while(!stack.isEmpty()){
                if(stack.peek() > find) {
                    sb.append(stack.pop() + " ");
                    findMax = true;
                    break;
                }else {
                    stack.pop();
                }
            }

            if(findMax != true) {
                sb.append(-1 + " ");
                findMax = false;
            }
            
            stack.clear();

        }

        sb.append(-1);

        System.out.println(sb);

    }
}

그렇다면 어떠한 방식으로 접근해야 하는가

스택에는 값이 아닌 그 값에 인덱스를 넣어주면서 지정한 원소에 오큰 수를 찾는 것이 아닌 그 수가 오큰수의 후보가 되는지 찾는 방식으로 풀면 될거 같다.

위와 같은 방식을 코드에 적용하면 다음과 같다.

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

class 오큰수 {

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

        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringBuilder sb = new StringBuilder();
        int N = Integer.parseInt(br.readLine());
        int[] A = new int[N];
        int[] result = new int[N];
        String[] str = br.readLine().split(" ");
        for(int i = 0; i < N; i++) {
            A[i] = Integer.parseInt(str[i]);
        }

        Stack<Integer> stack = new Stack<>();
        stack.push(0);
        for(int i = 1; i < N; i++) {
            while(!stack.isEmpty() && A[stack.peek()] < A[i]) {
                result[stack.pop()] = A[i];
            }
            stack.push(i);
        }

        while(!stack.isEmpty()) {
            result[stack.pop()] = -1;
        }

        for(int i = 0; i < result.length; i++){
            sb.append(result[i] + " ");
        }

        System.out.println(sb);
    }
}