Just Do IT!

파이썬에서 heapq 모듈 사용하기 본문

코딩테스트 준비/알고리즘 정리

파이썬에서 heapq 모듈 사용하기

MOON달 2024. 1. 26. 18:42
728x90

 

 

힙 자료구

 

특정한 규칙을 가지는 트리로, 최댓값과 최솟값을 찾는 연산을 빠르게 하기 위해 고안된 완전이진트리를 기본으로 한다.

 

min heap을 사용하면 원소들이 항상 정렬된 상태로 추가되고 삭제되며,min heap에서 가장 작은 값은 언제나 인덱스 0, 즉, 이진트리의 루트에 위치한다.

 

  • 최소 힙 : 부모 노드의 키값이 자식 노드의 키값보다 항상 작은 힙
  • 최대 힙 : 부모 노드의 키값이 자식 노드의 키값보다 항상 큰 힙

 

list와 tree의 장점을 모두 사용한다.

  1. 트리 구조 사용
    • 요소 삽입 및 최소 값의 제거 시 발생하는 요소의 검색 및 스왑의 횟수보다 일반적인 리스트를 사용할 때보다 적다.
  2. 리스트 사용
    • 완전 이진트리는 리스트로 코딩 가능
    • 리스트가 클래스보다 빠르다

 

 

 

 

 

 

파이썬 heapq 모듈

 

 heapq (priority queue) 알고리즘을 제공한다.

 

모든 부모 노드는 그의 자식 노드보다 값이 작거나 큰 이진트리(binary tree) 구조인데, 내부적으로는 인덱스 0에서 시작해 k번째 원소가 항상 자식 원소들(2k+1, 2k+2) 보다 작거나 같은 최소 힙의 형태로 정렬된다.   

heapq는 내장 모듈로 별도의 설치 작업 없이 바로 사용할 수 있다.

 

 

import heapq

 

간단하게 import 해서 사용할 수 있다.

 

 

 

 

 

 

 

heapify

 

이미 생성해둔 리스트가 있다면 heapify 함수를 통해 힙 자료형으로 변환할 수 있다.

 

프로그래머스 level 2 더 맵게 문제를 통해 예시를 알아보자.

import heapq

def solution(scoville, K):
    answer = 0
    heapq.heapify(scoville)

 

 

 

heappop

 

가장 작은 원소를 힙에서 제거함과 동시에 그를 결과값으로 return한다.

import heapq

def solution(scoville, K):
    answer = 0
    heapq.heapify(scoville)
    
    while scoville[0] < K:
        new_arr = heapq.heappop(scoville) + (heapq.heappop(scoville) * 2)

 

 

반대로 heappush 는 item 값을 heap으로 push한다.

 

 

 

 

 

자세한 것은 아래 링크를 통해 알 수 있다.

https://docs.python.org/ko/3/library/heapq.html

 

heapq — Heap queue algorithm

Source code: Lib/heapq.py This module provides an implementation of the heap queue algorithm, also known as the priority queue algorithm. Heaps are binary trees for which every parent node has a va...

docs.python.org

 

 

 

 

 

 

 

 

추가 예시 - 최대 힙 (출처: https://www.daleseo.com/python-heapq/)

 

 heapq 모듈은 최소 힙 기능으로 동작하기 때문에 최대 힙으로 활용하려면 응용을 해야 한다.

힙에 튜플(tuple)를 원소로 추가하거나 삭제하면 튜플 내에서 맨 앞에 있는 값을 기준으로 최소 힙이 구성되는 원리를 이용하는 것이다.

 

즉,

각 값에 대한 우선순위를 구한 후,

(우선 순위, 값) 구조의 튜플(tuple)을 힙에 추가하거나 삭제하면 된다.

 

from heapq import heappush, heappop

nums = [4, 1, 7, 3, 8, 5]
heap = []

for num in nums:
  heappush(heap, (-num, num))  # (우선 순위, 값)

while heap:
  print(heappop(heap)[1])  # index 1

 

출력 시,

8
7
5
4
3
1

 

이런 형식으로 출력 된다.

'코딩테스트 준비 > 알고리즘 정리' 카테고리의 다른 글

완전 탐색 (Brute Force Algorithm)  (0) 2024.02.02