[프로그래머스/Lv.2] 뒤에 있는 큰 수 찾기 | 스택 | 파이썬

문제

https://school.programmers.co.kr/learn/courses/30/lessons/154539

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

 

제한사항

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

 

입출력 예

입출력 예 #1
2의 뒷 큰수는 3입니다. 첫 번째 3의 뒷 큰수는 5입니다. 두 번째 3 또한 마찬가지입니다. 5는 뒷 큰수가 없으므로 -1입니다. 위 수들을 차례대로 배열에 담으면 [3, 5, 5, -1]이 됩니다.

입출력 예 #2
9는 뒷 큰수가 없으므로 -1입니다. 1의 뒷 큰수는 5이며, 5와 3의 뒷 큰수는 6입니다. 6과 2는 뒷 큰수가 없으므로 -1입니다. 위 수들을 차례대로 배열에 담으면 [-1, 5, 6, 6, -1, -1]이 됩니다.

 

나의 풀이 (실패)

def solution(numbers):
    answer = []
    
    for i, n in enumerate(numbers):
        if max(numbers[i:]) <= n:
            answer.append(-1)
        else:
            for n2 in numbers[i:]:
                if n2 > n:
                    answer.append(n2)
                    break
            
    return answer

시간 초과로 실패했다.

이중 for문으로 인한 시간복잡도를 줄여야 할 거 같다.

 

다른 사람의 풀이

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

스택에 `numbers`의 인덱스를 하나씩 추가하면서 

스택의 마지막 원소를 인덱스로 가지는 `numbers`의 값이 현재의 수보다 작은 경우

스택에서 `pop()`하고, `answer`의 해당 인덱스의 값을 현재의 수로 변경한다.