二叉堆

What · 是什么

二叉堆(Binary Heap)是一棵完全二叉树，每个节点的权值不小于(大根堆)或不大于(小根堆)其儿子。堆性质保证根即为最值，O(1)查询。它支持三种操作：插入(末尾+向上调整sift up)、删除最值(根与末尾交换+向下调整sift down)、增减某点权值(相应向上/向下调整)，各O(log n)。用数组h存储时h[i]的两个儿子是h[2i]与h[2i+1]，根为h[1]。建堆有O(n)的向下调整法(从叶子逐个向下调整, 理解成每次合并两个已调整好的堆)与O(n log n)的逐元素向上调整法。堆是堆排序、Dijkstra优先队列、哈夫曼树、对顶堆(动态求中位数/第k大)等的基础。Python内置heapq实现小根堆(取负实现大根堆)。

Why · 为什么

How · 实现

伪代码

# 大根堆, h[1..n], 根1
sift_up(x):
  while x>1 and h[x]>h[x//2]:
    swap(h[x], h[x//2]); x//=2
sift_down(x):
  while 2*x<=n:
    t=2*x
    if t+1<=n and h[t+1]>h[t]: t+=1
    if h[x]>=h[t]: break
    swap(h[x], h[t]); x=t
insert(v): h[++n]=v; sift_up(n)
pop(): res=h[1]; h[1]=h[n--]; sift_down(1)
build_linear(): for i in n//2..1: sift_down(i)   # O(n)

Python 实现

class MaxHeap:
    """大根堆(数组实现).
    h[0]占位, 有效元素1..n; 支持push/pop/top/size.
    """
    def __init__(self, data=None):
        self.h = [0]
        if data:
            self.h.extend(data)
            self._build()
        else:
            self.h = [0]

    def _sift_up(self, x):
        h = self.h
        while x > 1 and h[x] > h[x // 2]:
            h[x], h[x // 2] = h[x // 2], h[x]
            x //= 2

    def _sift_down(self, x):
        h = self.h
        n = len(h) - 1
        while 2 * x <= n:
            t = 2 * x
            if t + 1 <= n and h[t + 1] > h[t]:
                t += 1
            if h[x] >= h[t]:
                break
            h[x], h[t] = h[t], h[x]
            x = t

    def _build(self):
        for i in range((len(self.h) - 1) // 2, 0, -1):
            self._sift_down(i)

    def push(self, v):
        self.h.append(v)
        self._sift_up(len(self.h) - 1)

    def pop(self):
        if len(self.h) == 1:
            raise IndexError("pop from empty heap")
        top = self.h[1]
        self.h[1] = self.h[-1]
        self.h.pop()
        if len(self.h) > 1:
            self._sift_down(1)
        return top

    def top(self):
        return self.h[1]

    def __len__(self):
        return len(self.h) - 1


class MinHeap:
    """小根堆: 用负值存入MaxHeap."""
    def __init__(self, data=None):
        self._h = MaxHeap([-v for v in data] if data else None)

    def push(self, v):
        self._h.push(-v)

    def pop(self):
        return -self._h.pop()

    def top(self):
        return -self._h.top()

    def __len__(self):
        return len(self._h)


def heap_sort(arr):
    """用二叉堆排序."""
    h = MaxHeap(arr)
    res = []
    while len(h):
        res.append(h.pop())
    return res


class MedianStream:
    """对顶堆: 动态维护流的中位数.
    小根堆存较大半区, 大根堆存较小半区, 中位数=小根堆顶.
    """
    def __init__(self):
        self.small = MaxHeap()   # 较小半区(大根堆)
        self.large = MinHeap()   # 较大半区(小根堆)

    def add(self, v):
        if len(self.small) == 0 or v <= self.small.top():
            self.small.push(v)
        else:
            self.large.push(v)
        self._rebalance()

    def _rebalance(self):
        while len(self.small) > len(self.large) + 1:
            self.large.push(self.small.pop())
        while len(self.large) > len(self.small):
            self.small.push(self.large.pop())

    def median(self):
        if len(self.small) == 0:
            return None
        if (len(self.small) + len(self.large)) % 2 == 1:
            return self.small.top()
        return self.small.top()  # 偶数时输出较小中位数


if __name__ == "__main__":
    import random

    # 堆性质与push/pop
    h = MaxHeap()
    for x in [3, 1, 4, 1, 5, 9, 2, 6]:
        h.push(x)
    assert h.top() == 9
    popped = [h.pop() for _ in range(len(h))]
    assert popped == sorted(popped, reverse=True), popped

    # 线性建堆
    h2 = MaxHeap([5, 3, 8, 1, 9, 2])
    assert h2.top() == 9
    assert [h2.pop() for _ in range(len(h2))] == sorted([5, 3, 8, 1, 9, 2],
                                                        reverse=True)

    # 堆排序
    arr = [random.randint(-100, 100) for _ in range(50)]
    assert heap_sort(arr) == sorted(arr, reverse=True)

    # 小根堆
    mh = MinHeap()
    for x in [4, 2, 8, 1, 7]:
        mh.push(x)
    assert mh.top() == 1

    # 对顶堆中位数: 插入 [1,2,3,4,5]
    ms = MedianStream()
    for v in [1, 2, 3, 4, 5]:
        ms.add(v)
    assert ms.median() == 3
    ms.add(6)
    assert ms.median() == 3  # 偶数长度输出较小中位数

    print("binary-heap: ALL TESTS PASSED")

易错点

• 索引从0还是1开始要统一: 根为1则儿子2i/2i+1, 根为0则儿子2i+1/2i+2, 混用必出错
• 向下调整需比较两个儿子取较大者, 且先判右儿子存在(t+1<=n)
• 删除最值时先与末尾交换再pop, 不要直接删根(会破坏完全二叉树结构)
• 建堆用向下调整法才O(n), 逐元素插入是O(n log n)
• heapq是最小堆, 取最大堆要push(-x)并在pop时取负

例题

信源

• OI Wiki✓ https://oi-wiki.org/ds/binary-heap/
• 书籍✓ CPH ch4 p47