뒤에 있는 큰 수 찾기

문제 링크

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

문제 풀이

문제 설명

정수로 이루어진 배열 numbers가 있을 때, 자신보다 뒤에 있는 숫자 중 자신보다 크면서 가장 가까이 있는 수를 뒷 큰수라고 한다.

뒷 큰수를 차례로 담은 배열을 return하는 문제이다.

문제 접근 방식

배열의 최대 크기는 100,000이라, O(N^2)을 하게 될 경우, 시간초과가 발생했다. 이를 해결하기 위해 스택을 활용했다.

구현 방식은 아래와 같다.

for문을 돌며 스택에 해당 index를 추가한다. 단, 추가하기 전 index에 해당하는 수가 스택의 마지막 값보다 더 큰지 체크를 한다. 만약 index에 해당하는 수가 더 크다면, 이는 뒷 큰수에 해당하여 while문을 돌게 된다. while문은 스택을 뒤에서부터 비우며 스택이 다 비워지거나, 스택에 있는 값이 index에 해당하는 값보다 클 경우에 종료된다.

최종 코드

최종 코드는 아래와 같다.

def solution(numbers):
    answer = [-1] * len(numbers)
    stack = [0]
    for i in range(len(numbers)):
        # 뒤에 있는 수가 더 크다면 청산
        while stack and numbers[i] > numbers[stack[-1]]:
            answer[stack.pop()] = numbers[i]
        # 스택에 index 추가
        stack.append(i)
    return answer