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

mileh 2024. 11. 17. 00:46

문제

정수로 이루어진 배열 numbers가 있습니다. 배열 의 각 원소들에 대해 자신보다 뒤에 있는 숫자 중에서 자신보다 크면서 가장 가까이 있는 수를 뒷 큰수라고 합니다.

정수 배열 numbers가 매개변수로 주어질 때, 모든 원소에 대한 뒷 큰수들을 차례로 담은 배열을 return 하도록 solution 함수를 완성해주세요. 단, 뒷 큰수가 존재하지 않는 원소는 -1을 담습니다.

제한사항

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

배열이 주어졌을 때, 자신 뒤에 있는 요소들 중 첫 번째로 큰 수를 조회하는 문제입니다. [2,3,3,5]로 예를 들어 보겠습니다.

2의 오른쪽에 있는 요소들 중 첫번 째로 큰 수는 3입니다.
3의 오른쪽에 있는 요소들 중 첫번 째로 큰수는 5입니다. (자신보다는 큰 수 이어야 합니다.)
마찬가지로 3의 오른쪽에 있는 요소들 중 첫번째로 큰수는 5가 됩니다.
마지막 5의 경우 자신의 오른쪽에 있는 요소가 없기 때문에 -1입니다.

백준 사이트의 오큰수 문제와 같은 문제입니다. 이전에 코딩테스트 책으로 풀어봤던 해설지를 참고하여 해결하였습니다.

풀이

배열 numbers의 각 원소를 순회하면서 스택을 사용해 현재 원소보다 큰 값을 찾습니다.
만약 numbers[i]가 스택에 저장된 인덱스 st.peek()의 값보다 크면, 스택의 인덱스를 꺼내서 해당 위치에 numbers[i]를 저장합니다. 이 과정으로 answer[st.pop()]에 큰 값을 기록합니다.
스택에 인덱스 i를 추가해 다음 비교를 준비합니다.
반복문이 끝난 후에도 스택에 남아 있는 인덱스들은 오른쪽에 자신보다 큰 값이 없는 경우이므로 -1을 저장합니다.

파이썬의 경우 더 간단하게 풀 수 있어 다른 방식으로 진행했습니다.

배열 numbers의 각 원소를 순회하면서 스택을 사용해 현재 원소보다 큰 값을 찾습니다.
만약 numbers[i]가 스택에 저장된 인덱스의 값보다 크면, stack.pop()으로 스택의 인덱스를 꺼내 해당 위치에 numbers[i]를 저장합니다. 이는 그 인덱스 위치의 원소가 처음 만난 더 큰 수이기 때문입니다.
현재 인덱스 i를 스택에 추가하여 다음 비교를 준비합니다.

Java

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

public class Solution {
    public int[] solution(int[] numbers) {
        int[] answer = new int[numbers.length];      
        Stack<Integer> st = new Stack<>();
        st.push(0);
        
        for(int i = 1; i<numbers.length; i++) {
            while (!st.isEmpty() && numbers[st.peek()] < numbers[i]) {
                answer[st.pop()] = numbers[i];
            }
            st.push(i);
        }
        while (!st.empty()) {
            answer[st.pop()] = -1;
        }
        return answer;
    }
}

Python

def solution(numbers):
    stack = []
    answer = [-1] * len(numbers)

    for i in range(0,len(numbers)):
        while stack and numbers[stack[-1]] < numbers[i]:
            answer[stack.pop()] = numbers[i]
        stack.append(i)

    return answer