[프로그래머스/Java] 뒤에 있는 큰 수 찾기

문제 설명

https://school.programmers.co.kr/learn/courses/30/lessons/154539

코딩테스트 연습 - 뒤에 있는 큰 수 찾기

정수로 이루어진 배열 numbers가 있습니다. 배열의 각 원소들에 대해 자신보다 뒤에 있는 숫자 중에서 자신보다 크면서 가장 가까이 있는 수를 뒷 큰 수라고 합니다.
정수 배열 numbers가 매개변수로 주어질 때, 모든 원소에 대한 뒷 큰수들을 차례로 담은 배열을 return 하도록 solution 함수를 완성해 주세요. 단, 뒷 큰 수가 존재하지 않는 원소는 -1을 담습니다.

제한 사항

• 4 ≤ numbers의 길이 ≤ 1,000,000
  • 1 ≤ numbers[i] ≤ 1,000,000

입출력 예

numbers	result
[2, 3, 3, 5]	[3, 5, 5, -1]
[9, 1, 5, 3, 6, 2]	[-1, 5, 6, 6,-1, -1]

입출력 예 설명

입출력 예 #1

• 2의 뒷 큰 수는 3입니다. 첫 번째 3의 뒷 큰수는 5입니다. 두 번째 3 또한 마찬가지입니다.
  5는 뒷 큰 수가 없으므로 -1입니다.
  위 수들을 차례대로 배열에 담으면 [3, 5, 5, -1]이 됩니다.

입출력 예 #2

• 9는 뒷 큰수가 없으므로 -1입니다.
  1의 뒷 큰 수는 5이며, 5와 3의 뒷 큰수는 6입니다.
  6과 2는 뒷 큰 수가 없으므로 -1입니다.
  위 수들을 차례대로 배열에 담으면 [-1, 5, 6, 6, -1, -1]이 됩니다.

 

내가 작성한 코드

import java.util.*;

class Solution {
    public int[] solution(int[] numbers) {
        // 정답을 저장할 배열 생성
        int[] answer = new int[numbers.length];
        
        // 배열을 모두 -1로 채움
        Arrays.fill(answer, -1);
        
        // 인덱스를 저장하는 스택
        Deque<Integer> stack = new ArrayDeque<>();
        
        // 뒤에서부터 탐색
        for (int i = numbers.length - 1; i >= 0; i--) {
            int num = numbers[i];

            // 현재 숫자보다 작거나 같은 값은 다음 큰 수가 될 수 없으므로 제거
            while (!stack.isEmpty() && stack.peek() <= num) {
                stack.pop();
            }
            
            // 스택이 비어있지 않다면, 맨 위의 값이 '뒤에 있는 큰 수'
            if (!stack.isEmpty()) {
                answer[i] = stack.peek();
            }
            
            // 현재 값을 스택에 추가
            stack.push(num);
        }
        
        return answer;
    }
}

문제를 풀기 위해 생각한 방향

처음 문제를 봤을 때, 각 숫자에 대해 자신보다 뒤에 있는 숫자 중에서 처음으로 자신보다 큰 수를 찾아야 한다는 점을 파악했다. 따라서 가장 직관적으로는, 배열을 순회하면서 현재 숫자 이후의 원소들을 하나씩 비교하여 자신보다 큰 수를 찾으면 답에 저장하는 방식이 떠올랐다. 만약 끝까지 가도 큰 수가 없다면 -1을 저장하도록 했다. 이 방법은 구현이 단순하고 직관적이지만, 최악의 경우 모든 원소에 대해 뒤쪽을 끝까지 탐색해야 하므로 시간 복잡도가 O(N^2)이 되어 입력 크기가 커지면 효율성 테스트를 통과하기 어렵다. 실제로 기본 테스트는 통과했지만, 숨은 테스트에서는 시간 초과가 발생했다.

이 문제를 효율적으로 풀기 위해서는 불필요한 비교를 줄여야 했다. 그래서 배열을 뒤에서부터 보면서, 스택을 활용해 "뒤에 있는 큰 수" 후보들을 관리하는 방식을 생각했다. 구체적으로는 스택에 항상 내림차순으로 값이 유지되도록 하여, 현재 숫자보다 작거나 같은 값들은 스택에서 제거하고, 남아 있는 스택의 맨 위 값이 곧 현재 숫자의 "뒤에 있는 큰 수"가 된다. 만약 스택이 비어 있다면 뒤에 자신보다 큰 수가 없는 것이므로 -1로 처리하면 된다. 이 방식은 각 원소가 한 번 push, 한 번 pop 되므로 전체 시간 복잡도는 O(N)에 해결할 수 있어 효율성 테스트도 통과할 수 있다.

결국 이 문제는 단순한 완전 탐색 방식으로도 풀 수 있지만, 입력 크기에 따라 효율성이 중요한 상황에서는 단조 감소 스택을 활용하는 것이 최적의 풀이임을 알 수 있었다.