[Python] heapq, bisect, array

heapq: min-heap

heapq는 리스트를 min-heap 구조로 다루는 함수 모음. 클래스가 아니라 일반 list 위에서 동작.

기본

import heapq

h = []
heapq.heappush(h, 5)
heapq.heappush(h, 1)
heapq.heappush(h, 3)
heapq.heappush(h, 2)

print(heapq.heappop(h))    # 1 (최소)
print(heapq.heappop(h))    # 2
print(h)                    # [3, 5] (내부 표현, 정렬 X)

heappush/heappop 모두 O(log n). 단순 list로는 정렬에 O(n log n) 또는 최소 검색에 O(n).

리스트를 heap으로 변환

import heapq
​
xs = [3, 1, 4, 1, 5, 9, 2, 6]
heapq.heapify(xs)            # in-place, O(n)
print(xs)
​
# 작은 것부터 pop
while xs:
  print(heapq.heappop(xs), end=" ")

[1, 1, 2, 3, 5, 9, 4, 6]
1 1 2 3 4 5 6 9

heapify(xs)는 in-place O(n). 내부 배열은 heap 속성만 만족, 일반 정렬 아님.

nlargest / nsmallest

import heapq

heapq.nlargest(3, [1, 5, 2, 8, 3])     # [8, 5, 3]
heapq.nsmallest(3, [1, 5, 2, 8, 3])    # [1, 2, 3]

# key 함수
heapq.nlargest(2, students, key=lambda s: s.score)

전체를 정렬할 필요 없이 top-N만 가져온다. N이 작으면 매우 효율적 (O(n log k)).

max-heap

heapq는 min-heap만. max-heap이 필요하면 부호 뒤집기.

import heapq

h = []
heapq.heappush(h, -5)       # 음수로 저장
heapq.heappush(h, -1)
heapq.heappush(h, -3)

print(-heapq.heappop(h))    # 5 (최대)

튜플 키도 가능:

# (-priority, item) 형태
heapq.heappush(h, (-priority, item))
_, item = heapq.heappop(h)

우선순위 큐 with 동률

같은 우선순위에서 삽입 순서 보존하려면 카운터 추가.

import heapq
import itertools

counter = itertools.count()
pq = []

def push(priority, item):
    heapq.heappush(pq, (priority, next(counter), item))

def pop():
    priority, _, item = heapq.heappop(pq)
    return item

push(3, "task3")
push(1, "task1")
push(3, "task3.5")
print(pop())    # task1
print(pop())    # task3 (먼저 들어간 것)
print(pop())    # task3.5

객체끼리 비교가 불가능할 때 동률 시 TypeError 방지 효과도.

heappushpop / heapreplace

heapq.heappushpop(h, x)   # push 후 pop (한 번에 더 효율적)
heapq.heapreplace(h, x)   # pop 후 push (역순)

heapreplace는 비어 있으면 IndexError.

merge: 정렬된 시퀀스 병합

import heapq

a = [1, 3, 5, 7]
b = [2, 4, 6, 8]
c = [0, 9]

list(heapq.merge(a, b, c))    # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

각 입력이 이미 정렬되어 있다고 가정. iterator라 메모리 효율적.

bisect: 정렬된 리스트의 이진 탐색

import bisect

xs = [1, 3, 5, 7, 9]

print(bisect.bisect_left(xs, 5))    # 2 (== 5의 가장 왼쪽 인덱스)
print(bisect.bisect_right(xs, 5))   # 3 (== 5의 가장 오른쪽 이후)
print(bisect.bisect(xs, 5))         # 3 (bisect_right와 동일)

# 정렬 유지하며 삽입
bisect.insort(xs, 4)
print(xs)                            # [1, 3, 4, 5, 7, 9]

O(log n) 이진 탐색 + O(n) 삽입(리스트 이동). 작은 리스트엔 빠르고 명확.

응용: rank, sorted insert

import bisect

# 점수 → 등급
def grade(score, thresholds=[60, 70, 80, 90]):
    """thresholds 기준 등급 반환"""
    grades = "FDCBA"
    i = bisect.bisect_left(thresholds, score)
    return grades[i]    # 잘못 - inverted

# 더 흔한 패턴: 정렬 순위
def rank(arr, x):
    return bisect.bisect_left(sorted(arr), x)

key 인자 (3.10+)

import bisect

users = [{"age": 20}, {"age": 25}, {"age": 30}]
bisect.bisect_left(users, 26, key=lambda u: u["age"])    # 2

이전엔 키 시퀀스 별도 유지가 필요했음.

array: 컴팩트한 동질 배열

list는 PyObject 포인터 배열이라 int 하나에도 28바이트(+ 포인터 8바이트) 정도. array.array는 C 타입으로 저장 → 메모리 1/4 ~ 1/8.

from array import array

a = array("i", [1, 2, 3, 4, 5])    # signed int
a.append(6)
a.extend([7, 8])

print(len(a), a.itemsize, a.itemsize * len(a))

타입 코드:

list 대비 메모리 절감 효과는 큼. 그러나 산술 연산은 여전히 Python 수준. 수치 계산은 numpy가 압도적.

import array
import numpy as np

# array.array: 메모리 절감만
a = array.array("i", range(100000))

# numpy: 메모리 + 속도 (벡터 연산)
b = np.arange(100000)
b * 2 + 1    # 한 번에 벡터 연산

비교: 언제 무엇을 쓰나

sortedcontainers (third-party 추천)

pip install sortedcontainers

from sortedcontainers import SortedList, SortedDict, SortedSet

sl = SortedList([3, 1, 4, 1, 5])
sl.add(2)            # 정렬 유지 삽입 O(log n)
print(sl[0])         # 최소 O(1)
print(sl[-1])        # 최대 O(1)
print(sl.index(4))   # 위치
sl.remove(1)         # 삭제 O(log n)

표준 라이브러리엔 없지만 알고리즘 문제·실무에서 자주 유용. 내부적으로 list의 list로 구현 → 평균 매우 빠름.

자주 보는 패턴

top-k 스트리밍

import heapq

def top_k(stream, k):
    h = []
    for x in stream:
        if len(h) < k:
            heapq.heappush(h, x)
        elif x > h[0]:
            heapq.heapreplace(h, x)
    return sorted(h, reverse=True)

전체 스트림을 메모리에 안 들고 top-k만 유지. O(n log k).

k-way merge

import heapq

def kway_merge(*sorted_iters):
    return list(heapq.merge(*sorted_iters))

외부 정렬(external sort)에서 자주 등장하는 패턴.

Dijkstra 알고리즘

import heapq

def dijkstra(graph, start):
    dist = {start: 0}
    pq = [(0, start)]
    while pq:
        d, u = heapq.heappop(pq)
        if d > dist.get(u, float('inf')): continue
        for v, w in graph[u].items():
            nd = d + w
            if nd < dist.get(v, float('inf')):
                dist[v] = nd
                heapq.heappush(pq, (nd, v))
    return dist

우선순위 큐의 대표적 응용.