[PS BOJ 17298] - 오큰수

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

리뷰하게 된 이유

계속해서 시도를 했던 문제이지만, 결국 스스로 답을 찾아내지는 못했던 문제라 리뷰를 하면서 정리하고자 생각하게 됐다.

STEP 2. 구현

초기 아이디어

예시 입력을 참고해서 설명해보겠다.

[3, 5, 2, 7] → NGE = [5, 7, 7, -1]

오큰수는 큰 수 중에서 가장 왼쪽에 있는 수를 의미한다.

따라서, 3보다 크고 가장 왼쪽에 있는 수인 5가 3의 NGE가 된다.

5를 보자면 5는 2보다 크므로 무시되고, 7이 5의 NGE가 된다.

마찬가지로 7은 2의 NGE가 된다.

7은 자기 자신보다 큰 수가 없으므로 NGE = -1이 된다.

처음에 이 문제를 풀려고 착안했었던 부분은 첫번째 인덱스를 제외한 부분을 Queue로 세팅한 뒤에

3 → [5, 2, 7] queue 비교

NGE = 5

5 → [5, 2, 7] queue 비교

NGE 해당 값이 없으므로 dequeue

5 → [2, 7] queue 비교

NGE 해당 값이 없으므로 dequeue

5 → [7]

NGE = [5, 7]

2 → [7]

NGE = [5, 7, 7]

7 → [7] 마지막 인덱스이며 같으므로 NGE = -1

이런식으로 풀고자 했다. 하지만, 반례가 존재했다.

9 5 4 8을 넣는다 가정하자.

만약, 처음 인덱스가 제일 큰 값인 경우에는 queue의 값이 다 사라지게 된다.

이 부분을 해결하지 못했었다.

그렇다면 어떻게 이런 문제를 해결하고 문제 접근을 할 수 있을까? 고민을 해봐도 딱히 답이 안떠올랐다.

많은 사람들이 푼 답은 바로 ‘스택’ 이었다.

스택을 활용한 아이디어

스택을 활용한 아이디어는 정답 배열(nge)과 스택이 필요하다.

아이디어는 다음과 같다.

예시 : 3 5 2 7 (inputArr이라 명명)

i = 0 일때, 초기에 오큰수를 알아낼 수 없으니 스택에 예시 배열의 인덱스를 추가한다.

Stack : inputArr[0]

i = 1 일때, 스택이 비어있지 않으므로, stack의 top과 inputArr[i] 비교한다.

Stack : inputArr[0] = 3

inputArr[1] = 5

이때, inputArr[i] > stack의 top이면 스택의 탑을 뽑아내서 정답 배열에 저장 후 inputArr[i]를 다시 stack에 push한다.

→ if(inputArr[i] > stack.peek()) then 
    nge[stack.pop()] = inputArr[i]; → nge의 index에 오큰수를 넣어주는 개념
    stack.push(inputArr[i]);

위의 식에 의거해서 nge[0] = 5가 될 것이다.

그렇다면 inputArr[i] < stack의 top인 경우는 어떻게 처리해야될까?

i = 2 : 비교될 대상 숫자 2

Stack : inputArr[1] = 5

nge = [5, -1, -1, -1]

답은 아주 단순하게 이 경우에는 해당 값을 stack에 push한다.

Stack : [ inputArr[1], inputArr[2] ]

마지막으로 해당 두 개의 과정을 반복하면 된다.

i = 3 : 비교될 대상 숫자 7

Stack : [ inputArr[1] ] → inputArr[2] pop

nge[2] = inputArr[3]

nge = [5, -1, 7, -1]

Stack : [] → inputArr[1] pop

nge[1] = inputArr[3]

nge = [5, 7, 7, -1]

그리고 마지막 인덱스는 항상 -1이니 애초에 nge를 -1로 초기화하면 위의 과정으로 값을 도출 할 수 있다.

그렇다면? 위에서 해결 못한 반례를 해결 할 수 있을까?

[9, 5, 4, 8]

Stack : [0]

nge = [-1, -1, -1, -1]

i = 0 (9)

sequence[i] = 9 == sequence[stack.peek()] = 9

Stack : [0, 0]

nge = [-1, -1, -1, -1]

i = 1 (5)

sequence[i] = 5 < sequence[stack.peek()] = 9

Stack : [0, 0, 1]

nge = [-1, -1, -1, -1]

i = 2 (4)

sequence[i] = 4 < sequence[stack.peek()] = 5

Stack : [0, 0, 1, 2]

nge = [-1, -1, -1, -1]

i = 3 (8)

sequence[i] = 8 > sequence[stack.peek()] = 4

Stack : [0, 0, 1]

nge = [-1, -1, 8, -1]

sequence[i] = 8 > sequence[stack.peek()] = 5

Stack : [0, 0]

nge = [-1, 8, 8, -1]

sequence[i] = 8 < sequnece[stack.peek()] = 9

Stack : [0, 0, 3]

nge = [-1, 8, 8, -1]

로 옳은 정답이 나옴을 알 수 있다.

STEP 3. 구현 소스코드

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

public class Main {

    private final static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    private final static BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
    private final static Stack<Integer> stack = new Stack<>();
    
    public static void main (String[] args) throws IOException {
        int[] sequence = init();
        int[] nge = new int[sequence.length];
        
		// 초기 배열을 -1로 초기화한다. 
        Arrays.fill(nge, -1); 
        
        solve(sequence, nge);
        print(sequence.length, nge);
    }

    public static int[] init() throws IOException {

        int numCnt = Integer.parseInt(br.readLine());
        int[] sequence = new int[numCnt];

        String[] inputNumbers = br.readLine().split(" ");

        strToInt(sequence, inputNumbers);
        return sequence;
    }

    private static void strToInt(int[] sequence, String[] str) {
        for(int i = 0; i < str.length; i++){
            sequence[i] = Integer.parseInt(str[i]);
        }
    }

    public static void solve(int[] sequence, int[] nge) {
		// 초기 스택 세팅
		stack.push(0);
		
		for(int i = 0; i < sequence.length; i++) {
			while(!stack.isEmpty()){
                // [3, 5, 2, 7] i = 0인 경우 sequence[i] = 3 / sequence[stack.peek()] = 3
                // i = 1인 경우 sequence[i] = 5 / sequence[stack.peek()] = 3
				if(sequence[i] > sequence[stack.peek()]) {
                    // nge[stack.peek() = 0] = 5  
					nge[stack.pop()] = sequence[i];
				} else {
                    // i = 2인 경우 sequence[i] = 2 / sequence[stack.peek()] = 5
                    // stack = [5, 2] 
					stack.add(i);
					break;
				}
			}	
			// 스택이 다 빌 경우에는 해당 값을 넣어주면 된다. 
			if(stack.isEmpty()){
				stack.add(i);
			}
		}
    }
    
    private static void print(int size, int[] nge) throws IOException {
        for(int i = 0; i < size; i++) {
            bw.write(nge[i] + " ");
        }
        bw.flush();
        bw.close();
    }

}

STEP 4. 결론

스택의 인덱스로 처리하는 방식은 생각하지 못 했었다.

앞으로 더 많은 문제를 풀어봐야겠다는 생각을 가지게 되는 것 같다.

댓글남기기