[프로그래머스] 뒤에 있는 큰 수 찾기 (Python) - heapq

문제 설명

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

---

제한사항

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

---

입출력 예

 

numbers	result
[2, 3, 3, 5]	[3, 5, 5, -1]
[9, 1, 5, 3, 6, 2]	[-1, 5, 6, 6, -1, -1]

---

입출력 예 설명

입출력 예 #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]이 됩니다.

---

첫번째 풀이 (오답)

from collections import deque

def solution(numbers):
    answer = [-1 for _ in numbers]
    q = deque(numbers)
    
    now_i = 0
    no_answer = deque() # 아직 답을 찾지 못한 인덱스들
    while q:
        now = q.popleft()
        
        new_no_answer = deque()
        while no_answer:
            i = no_answer.pop()
            if answer[i] < 0: 
                if numbers[i] < now:
                    answer[i] = now
                else: # 답을 못 찾은 경우 다시 노답리스트에 추가
                    new_no_answer.append(i)
                
        new_no_answer.append(now_i)
        no_answer = new_no_answer
        now_i += 1
        
    return answer

큐와 인덱스를 사용한 풀이.

답을 아직 찾지 못한 인덱스를 저장하는 방법을 사용했지만... 시간초과로 인해서 6개의 테케를 통과 못했다.

18번 테케 시간 보면 ㅋㅋㅋㅋㅋㅋ 9561 .....

실활까 ?...

큐 쓰면 무조건 잘 될 줄 알았는데 ..(?)

 

.. 결국 통과를 못하고 검색을 통해 얻은 지식 = 힙을 사용해봐라.

생각해보니 인덱스를 사용하는 것에 너무 집착한 것 같아서

그냥 숫자를 사용하고 힙을 사용해 보기로 했다.

 

힙 중에서도 최소힙!

다른말로는 우선순위 큐인데 숫자가 작은 걸 우선순위로 갖는 큐!

를 사용해보자.

 

두번째 풀이 (정답)

from heapq import heappush, heappop

def solution(numbers):
    answer = []
    
    heap = []
    for idx, num in enumerate(numbers):
        heappush(heap, [num,idx])
        while heap and heap[0][0] < num:
            heappop(heap)
        print(heap)
    return answer

사실 heap을 사용할 때 배열을 넣어도 되는 줄 몰랐다.

배열의 첫번째 원소를 기준으로 정렬해주는구나 .. 파이썬은 역시 똑똑해 ....

구현은 다 했는데 고민하다가 배열을 한번 넣어봤는데 잘 되어서 이렇게 또 하나 배워간다.

 

역시 파이썬은 만능이야 ...

자바로 문제 풀던 시절이 하나도 기억이 안나고 자바는 아주 바보같다는 생각 뿐 🤣

 

하지만 역시 가장 바보는 나야 ... 언제쯤 검색 안하고 풀 수 있을까 .........?,,.,,,