이중우선순위 큐 (힙) - PYTHON

문제 설명

이중 우선순위 큐는 다음 연산을 할 수 있는 자료구조를 말합니다.

명령어 수신 탑(높이)
I 숫자 큐에 주어진 숫자를 삽입합니다.
D 1 큐에서 최댓값을 삭제합니다.
D -1 큐에서 최솟값을 삭제합니다.
이중 우선순위 큐가 할 연산 operations가 매개변수로 주어질 때, 모든 연산을 처리한 후 큐가 비어있으면 [0,0] 비어있지 않으면 [최댓값, 최솟값]을 return 하도록 solution 함수를 구현해주세요.

제한사항

operations는 길이가 1 이상 1,000,000 이하인 문자열 배열입니다.
operations의 원소는 큐가 수행할 연산을 나타냅니다.
원소는 “명령어 데이터” 형식으로 주어집니다.- 최댓값/최솟값을 삭제하는 연산에서 최댓값/최솟값이 둘 이상인 경우, 하나만 삭제합니다.
빈 큐에 데이터를 삭제하라는 연산이 주어질 경우, 해당 연산은 무시합니다.

입출력 예

operations return
[I 16,D 1][0,0]
[I 7,I 5,I -5,D -1][7,5]
입출력 예 설명
16을 삽입 후 최댓값을 삭제합니다. 비어있으므로 [0,0]을 반환합니다.
7,5,-5를 삽입 후 최솟값을 삭제합니다. 최대값 7, 최소값 5를 반환합니다.

풀이

최대 , 최소에 대한 문제는 sort를 사용해서 풀 수 있지만, 이전의 문제 풀이를 통해 sort보다 heapq을 사용한 최대, 최소 찾는 게 더 빠르다는 것을 알게 되었다. 좀 더 익숙해 지기 위해 sort보다는 heap을 사용하도록 하겠다.

알고리즘

• 최대, 최소 힙에 각각 값을 저장.
• D num 의 경우 ,
  - 최대 값 삭제 후 min heap 재 할당
• D -num 의 경우,
  - 최소 값 삭제 후 max heap 재 할당
• 각각 길이가 0 이면 0을 answer에 넣어줌. 아니면 최대 최소 값 넣어줌

import heapq

def changeHeap(heap) :
    res = []
    for h in heap :
        heapq.heappush(res, -h)
    return res

def solution(operations):
    answer = []
    min_heap = []
    max_heap = []
    
    for op in operations :
        cmd , num = op.split(' ')
        num = int(num)
        
        if cmd == 'I' :
            heapq.heappush(min_heap,num)
            heapq.heappush(max_heap,-num)
        else :
            try:
                if num == 1 :
                    heapq.heappop(max_heap)
                    min_heap = changeHeap(max_heap)
                else :
                    heapq.heappop(min_heap)
                    max_heap = changeHeap(min_heap)
            except:
                continue
    if len(max_heap) != 0:
        answer.append(-max_heap[0])
    else:
        answer.append(0)
        
    if len(min_heap) != 0:
        answer.append(min_heap[0])
    else:
        answer.append(0)
    
    return answer

그러나 다른 사람의 코드를 보는 도중에 heap을 사용한 더 간단한 코드를 보았다.

import heapq

def solution(operations):
    h = []

    for i in operations:
        a, b = i.split()
        if a == 'I':
            heapq.heappush(h, int(b))
        else:
            if len(h) > 0:
                if b == '1':
                    h.pop(h.index(heapq.nlargest(1, h)[0]))
                else:
                    heapq.heappop(h)

    if len(h) == 0:
        return [0, 0]
    else:
        return [heapq.nlargest(1, h)[0], h[0]]

이전의 코드 처럼 최소,최대 힙을 각자 설정하지 않아도 되는 장점이 있다.
이 코드의 핵심은 h.pop(h.index(heapq.nlargest(1, h)[0]))이 부분인데,
nlargest(n, heap) 함수는 n개의 가장 큰 값들로 이루어진 리스트를 반환한다. 즉, 첫번째 인자에 찾고싶은 최대값의 개수를 넣어주면 된다! 이로서 최대 값을 쉽게 찾을 수 있고 코드도 짧고 시간도 짧게 걸리는 것을 확인했다🤘