뒤에 있는 큰 수 찾기

문제 출처: https://school.programmers.co.kr/learn/courses/30/lessons/154539

이번 문제는 각 요소의 뒤에 있는 요소들 중에서 가장 가까운 큰 수를 찾는 문제입니다. 처음에는 다음과 같이 생각했습니다.

1. 각 요소를 순회한다.
2. 뒤 요소 중에서 값이 더 큰 요소가 있다면 answer 리스트에 추가하고 반복문을 빠져나온다.
3. 마지막 요소는 무조건 -1이므로 answer 리스트에 추가한 뒤 반환한다.

그래서 코드를 이렇게 작성했습니다.

def solution(numbers):
    answer = []
    for i in range(0, len(numbers) - 1):
        init_num = numbers[i]
        is_large = False
        
        for j in range(i+1, len(numbers)):
            if init_num < numbers[j]:
                answer.append(numbers[j])
                is_large = True
                break
        if not is_large:
            answer.append(-1)
        
    answer.append(-1)
    return answer

그런데 답안을 제출하니 시간 초과(Time Limit Exceeded)가 발생했습니다. 시간 복잡도가 O(n^2)여서 그런 것 같습니다.

한참 고민한 후에 스택을 사용하기로 했습니다. 스택 안에는 아직 자기보다 큰 수를 만나지 못한 인덱스를 넣어 둡니다.

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

동작 원리는 다음과 같습니다.

1. 먼저 모든 값을 -1로 초기화합니다.
2. 스택이 비어 있지 않다면, 스택이 빌 때까지 현재 배열의 숫자와 비교합니다.
3. 마지막에 현재 인덱스를 스택에 넣습니다.

이렇게 해서 테스트를 통과했습니다.