[프로그래머스] 뒤에 있는 큰 수 찾기 (파이썬)

문제 출처: 프로그래머스

프로그래머스

 

문제 설명

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

출처: 프로그래머스

---

이 문제는 이미 백준에서 한번 풀어본 적이 있는 문제였다. (=> 백준 오큰수)

 

예시를 맞추려고만 하면 아주 쉬운 문제이긴 하나, 시간 복잡도를 고려해서 정답을 맞추려면 조금 생각해야 하는 문제였다.

 

정답은 "자료구조" 안에 있다.

 

<정답 코드>

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

스택을 이용해 문제를 해결했다.

접근은 아래와 같다.

 

1. 정답을 기록할 answer 배열에 numbers의 길이만큼 -1을 입력한다.
2. number를 순회하되, stack을 이용해 순회하도록 작성한다.
   1. 순회 도중에 stack이 비어있을 경우도 있으므로 [0]으로 초기화해 빈 배열인지 점검할 수 있도록 한다.
   2. numbers[i]보다 작고, stack에 index가 기록되어있는 배열들은 전부 numbers[i]로 기록한다.

 

원리 자체는 쉬우나, 빨리빨리 떠올리기에는 연습이 필요한 문제인 것 같다.