[프로그래머스] Lv.2 뒤에 있는 큰 수 찾기 (Python / Java)

https://school.programmers.co.kr/learn/courses/30/lessons/154539?language=java 

프로그래머스

문제 설명

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

제한사항

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

입출력 예

 

---

문제 풀이

처음에는 이중 for문으로 문제를 풀었었는데 시간 초과가 떠서ㅠㅠ stack을 통해 문제를 풀었다.

1. 정답을 저장 할 answer 배열 생성하고, 초기값으로 모든 원소 -1로 설정
2. numbers 배열 순회
   • stack이 비어 있지 않고, 스택의 맨 위에 저장된 인덱스의 값이 현재 값보다 작다면
     • 현재 값이 스택의 맨 위에 저장된 인덱스의 값보다 크다는 것이므로 해당 값의 뒷 큰수로 설정
     • stack의 맨 위에 저장한 인덱스를 꺼내 answer 배열의 해당 위치에 현재 값(num) 할당
   • 현재 값의 인덱스를 스택에 추가
3. 결과 (answer 배열) 반환

 

[Python]

def solution(numbers):
    
    answer = [-1] * len(numbers)
    stack = []
    
    for idx, num in enumerate(numbers):
        while stack and numbers[stack[-1]] < num:
            answer[stack.pop()] = num      
        stack.append(idx)
        
    return answer

 

[Java]

import java.util.Stack;

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