[python] 뒤에 있는 큰 수 찾기

인증사진(2024-08-11)

문제

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

 

답

def solution(numbers):
    n = len(numbers)
    answer = [-1] * n  # 결과를 담을 배열, 기본값은 -1로 초기화
    stack = []  # 스택을 사용하여 뒤의 큰 수를 찾기
    
    for i in range(n - 1, -1, -1):
        # 현재 원소보다 작거나 같은 값들은 스택에서 제거
        while stack and stack[-1] <= numbers[i]:
            stack.pop()
        
        # 스택이 비지 않았다면 스택의 top이 현재 원소의 뒷 큰 수
        if stack:
            answer[i] = stack[-1]
        
        # 현재 원소를 스택에 추가
        stack.append(numbers[i])
    
    return answer

 

추가공부