heapq в Python и приёмы работы с максимальной кучей

Обзор модуля heapq в Python

import heapq

# heapq operates on plain Python lists
heap = []
heapq.heappush(heap, 5)
heapq.heappush(heap, 2)
heapq.heappush(heap, 8)
heapq.heappush(heap, 1)

print('Heap array:', heap)          # internal array (not sorted!)
print('Peek min:', heap[0])         # O(1) min access
print('Pop min:', heapq.heappop(heap))  # 1
print('Next min:', heap[0])         # 2

# heapify: turn any list into a heap in O(n)
data = [9, 4, 7, 1, 3, 6, 2]
heapq.heapify(data)
print('Heapified:', data, '| min:', data[0])

Макс-куча с помощью инвертирования значений

import heapq

max_heap = []
for val in [5, 1, 8, 3, 9, 2]:
    heapq.heappush(max_heap, -val)  # negate on push

print('Max-heap internal:', max_heap)  # all negated

# Pop in descending order:
results = []
while max_heap:
    results.append(-heapq.heappop(max_heap))  # negate on pop
print('Sorted descending:', results)  # [9, 8, 5, 3, 2, 1]

# Common pattern: top-k largest
data = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5]
k = 3
heap = []
for x in data:
    heapq.heappush(heap, -x)
print('Top', k, ':', [-heapq.heappop(heap) for _ in range(k)])

heapq.nlargest и nsmallest

import heapq

data = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 7]

# Top 3 largest:
print(heapq.nlargest(3, data))   # [9, 8, 7]
# Top 3 smallest:
print(heapq.nsmallest(3, data))  # [1, 1, 2]

# With a key function:
words = ['banana', 'apple', 'cherry', 'date', 'elderberry']
print(heapq.nlargest(2, words, key=len))   # ['elderberry', 'banana']
print(heapq.nsmallest(2, words, key=len))  # ['date', 'apple']

# Note: when k ~ n, use sorted() instead:
# sorted(data)[-k:]  or  sorted(data, reverse=True)[:k]

Куча с кортежами для сложных ключей

import heapq
import itertools

# Pattern: (priority, counter, item)
# Counter ensures unique tiebreaker, avoids comparing items
counter = itertools.count()
heap = []

def push_task(priority, task):
    heapq.heappush(heap, (priority, next(counter), task))

push_task(3, 'low priority task')
push_task(1, 'high priority task')
push_task(2, 'medium priority task')
push_task(1, 'another high priority')

while heap:
    pri, cnt, task = heapq.heappop(heap)
    print(f'P{pri}: {task}')
# Output in priority order: P1, P1, P2, P3

heapq.merge: объединение отсортированных итерируемых объектов

import heapq

# Merge multiple sorted lists efficiently
sorted_lists = [
    [1, 5, 9],
    [2, 6, 8],
    [3, 4, 7]
]

# heapq.merge takes sorted iterables and returns a merged sorted iterator
merged = list(heapq.merge(*sorted_lists))
print('Merged:', merged)  # [1, 2, 3, 4, 5, 6, 7, 8, 9]

# The k-way merge manually (educational version):
def merge_k_sorted(lists):
    heap = []
    for i, lst in enumerate(lists):
        if lst:
            heapq.heappush(heap, (lst[0], i, 0))
    result = []
    while heap:
        val, list_idx, elem_idx = heapq.heappop(heap)
        result.append(val)
        if elem_idx + 1 < len(lists[list_idx]):
            heapq.heappush(heap, (lists[list_idx][elem_idx+1], list_idx, elem_idx+1))
    return result

print('Manual k-way:', merge_k_sorted(sorted_lists))

Шаблон отложенного удаления для куч

import heapq

class LazyHeap:
    def __init__(self):
        self._heap = []
        self._removed = set()

    def push(self, task):
        heapq.heappush(self._heap, task)

    def remove(self, task):
        self._removed.add(task)  # mark as removed

    def pop(self):
        while self._heap:
            task = heapq.heappop(self._heap)
            if task not in self._removed:
                return task
        return None

lh = LazyHeap()
for t in [5, 1, 8, 3, 2]:
    lh.push(t)
lh.remove(1)  # 'delete' 1 lazily
lh.remove(8)  # 'delete' 8 lazily
results = [lh.pop() for _ in range(3)]
print(results)  # [2, 3, 5] -- 1 and 8 skipped

k-й по величине элемент в потоке

import heapq

class KthLargest:
    def __init__(self, k, nums):
        self.k = k
        self.heap = []
        for num in nums:
            self.add(num)

    def add(self, val):
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)  # remove smallest
        return self.heap[0]  # kth largest = root of min-heap

# k=3, initial=[4,5,8,2]
kl = KthLargest(3, [4, 5, 8, 2])
print(kl.add(3))   # 4 (top 3: 8,5,4 -- kth=4)
print(kl.add(5))   # 5 (top 3: 8,5,5 -- kth=5)
print(kl.add(10))  # 5 (top 3: 10,8,5 -- kth=5)
print(kl.add(9))   # 8 (top 3: 10,9,8 -- kth=8)

Найдите k пар с наименьшей суммой

import heapq

def k_smallest_pairs(nums1, nums2, k):
    if not nums1 or not nums2:
        return []
    heap = []
    # Initialize with pairs (nums1[0], nums2[j])
    for j in range(min(k, len(nums2))):
        heapq.heappush(heap, (nums1[0] + nums2[j], 0, j))
    result = []
    while heap and len(result) < k:
        total, i, j = heapq.heappop(heap)
        result.append([nums1[i], nums2[j]])
        if i + 1 < len(nums1):
            heapq.heappush(heap, (nums1[i+1] + nums2[j], i+1, j))
    return result

print(k_smallest_pairs([1,7,11], [2,4,6], 3))
# [[1,2], [1,4], [1,6]]

Планировщик задач с макс-кучей

import heapq
from collections import Counter

def least_interval(tasks, n):
    freq = Counter(tasks)
    heap = [-f for f in freq.values()]  # max-heap (negated)
    heapq.heapify(heap)
    time = 0
    while heap:
        cycle = n + 1
        temp = []
        for _ in range(cycle):
            if heap:
                temp.append(heapq.heappop(heap))
        for f in temp:
            if f + 1 < 0:  # still tasks remaining
                heapq.heappush(heap, f + 1)
        # Add full cycle or remaining tasks if queue empty
        time += cycle if heap else len(temp)
    return time

print(least_interval(['A','A','A','B','B','B'], 2))  # 8
print(least_interval(['A','A','A','B','B','B'], 0))  # 6

Куча в алгоритме Дейкстры

import heapq

def dijkstra(graph, start):
    dist = {node: float('inf') for node in graph}
    dist[start] = 0
    heap = [(0, start)]  # (distance, node)
    while heap:
        d, u = heapq.heappop(heap)
        if d > dist[u]:   # stale entry, skip
            continue
        for v, w in graph[u]:
            if dist[u] + w < dist[v]:
                dist[v] = dist[u] + w
                heapq.heappush(heap, (dist[v], v))
    return dist

graph = {
    'A': [('B', 4), ('C', 1)],
    'B': [('D', 1)],
    'C': [('B', 2), ('D', 5)],
    'D': []
}
print(dijkstra(graph, 'A'))  # {'A':0,'B':3,'C':1,'D':4}

Перестройка строки с помощью макс-кучи

import heapq
from collections import Counter

def reorganize_string(s):
    freq = Counter(s)
    heap = [(-f, c) for c, f in freq.items()]
    heapq.heapify(heap)
    result = []
    prev_freq, prev_char = 0, ''
    while heap:
        freq, char = heapq.heappop(heap)
        result.append(char)
        # Push back the previous character if still remaining
        if prev_freq < 0:
            heapq.heappush(heap, (prev_freq, prev_char))
        prev_freq, prev_char = freq + 1, char  # decrement freq (less negative)
    result_str = ''.join(result)
    # Verify no adjacent duplicates
    return result_str if len(result_str) == len(s) else ''

print(reorganize_string('aab'))   # 'aba'
print(reorganize_string('aaab'))  # '' (impossible)

Быстрая проверка

Итоги урока