[Programmers] LV2. 뒤에 있는 큰 수 찾기

문제

---

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

 

제한사항 

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

 

풀이

---

가장 쉽게 생각할 수 있는 방법은 반복문을 돌면서 현재 수 뒤에 있는 숫자들을 모두 돌면서 자기 자신보다 큰 수가 나오면 멈추는 것이다. 근데 이렇게 되면, 예를 들어서 [1,1,1,1,1,1,1,.......,999]이라면 결국 O(N^2)이 되면서 시간초과가 날 것이다(바로 접니다).

 

그래서 해당 문제의 접근법은 Stack 자료구조를 사용하는 것이다. 반복문을 돌면서 나오는 수가 Stack에 들어있는 수보다 크면 가장 가까운 수를 찾을 수 있다. 그리고 Stack에서는 pop 해준다. 이거를 반복해서 수행하는 것이다.

 

def solution(numbers):
    answers = [-1] * len(numbers) # 아무것도 없으면 -1
    stack = [] # 인덱스를 넣을 stack
    for idx, target in enumerate(numbers):
        while stack and numbers[stack[-1]] < target: # 만약 현재값이 더 크면?
            answers[stack.pop()] = target # 조건을 충족했으니 stack에서 빼주자
        stack.append(idx) 
        
    return answers