二叉堆相关知识点整理(Python实现)

二叉堆简介

做leetcode题目，解题说用二叉堆效率高，因此就有了这篇博客记录学习过程。二叉堆是一颗完全二叉树，满足其子节点总是小于父亲节点的特性（小根堆），大根堆则相反。python里可以使用heapq这个内置库操作堆，默认实现的是下根堆，如果在插入数据时，加上负号，取出时也加上负号，那么虽然存储的时候是按照小根堆的规则来，但是可以达到大根堆的效果

公式定理

一般使用数组来表示二叉堆，数组下标间有如下关系（具体推到过程，这里先挖个坑，下次回来填）：

如果从下标从0开始存储，则下标为i的结点的主要关系为：
双亲：（(i-1)//2）
左孩子：2i+1
右孩子：2i+2

操作：

建堆

建堆就是把一个数组变为符合小根堆定义的数组，需要一个起点，这个起点就是最后一个非叶子节点，这里有有一个公式：最后一个非叶子节点下标=n//2-1 (n为数组长度，下标从0开始)。
那么我们的思路就是：从最后一个非叶子节点开始，如果其值大于子节点的值，则交换。依次循环，直到根节点.

插入

在已经符合小根堆定义的情况下，再插入一个元素，使其再次符合最小堆的定义。我们有以下思路：新插入的值，可能会不符合小根堆定义，所以需要把新插入的值，与上级交换，简称上浮，直到堆顶

取值

我们要取数组中最小的数，其时间复杂度是O(1)，因为小根堆的堆顶就是。取完值，我们要把数组末尾最后一个元素移动到堆顶后，再使其符合小根堆定义，简称下沉。我们有以下思路：从堆顶开始，判断其是否大于其最小的子节点，如果是，则进行交换，直到堆尾。

源码参考

def down(arr, parent_index):
    """
    构建最小堆...
    和其叶子节点对比交换...所以是不断下沉，直到末尾
    :return:
    """
    temp = arr[parent_index]
    end_index = len(arr)
    child_index = 2 * parent_index + 1
    while child_index < end_index:
        if child_index + 1 < end_index and arr[child_index + 1] < arr[child_index]:
            child_index += 1
        # 由于是单向替换，要使用temp不应该使用arr[parent]
        if temp < arr[child_index]:
            break
        # 交换
        arr[parent_index] = arr[child_index]
        parent_index = child_index
        child_index = 2 * parent_index + 1
    arr[parent_index] = temp


def build_heap(arr):
    """
    从最后一个非叶子结点开始，直到根顶，依次使其满足最小堆定义
    最后一个非叶子节点满足定义：i = n/2 - 1 (n为数组长度)
    :return:
    """
    last_no_leaf_node = len(arr) // 2 - 1
    for i in range(last_no_leaf_node, -1, -1):
        down(arr, i)


def up(arr, start_index, curr_index):
    """
    把一个数不断上浮，直到堆顶
    :param curr_index:
    :param arr:
    :param start_index:
    :return:
    """
    new_item = arr[curr_index]
    while curr_index > start_index:
        parent_index = (curr_index - 1) // 2
        if arr[parent_index] < new_item:
            break
        arr[curr_index] = arr[parent_index]
        curr_index = parent_index
    arr[curr_index] = new_item


def heappush(arr, num):
    """
    往堆尾插入一个数，这个数需要上浮到合适位置...
    :param arr:
    :param num:
    :return:
    """
    arr.append(num)
    up(arr, 0, len(arr) - 1)


def heappop(arr):
    """
    返回堆顶元素
    :param arr:
    :return:
    """
    last_item = arr.pop()
    if arr:
        return_item = arr[0]
        arr[0] = last_item
        down(arr, 0)
        return return_item
    return last_item


if __name__ == '__main__':
    a = [4, 3, 5, 1, 0]
    print(a)
    # 建堆操作
    build_heap(a)
    print(a)
    # 往堆里插入一个数
    heappush(a, 2)
    print(a)
    # 返回堆顶元素
    result = heappop(a)
    print(result)
    print(a)
    
	# [4, 3, 5, 1, 0]
	# [0, 1, 5, 4, 3]
	# [0, 1, 2, 4, 3, 5]
	# 0
	# [1, 3, 2, 4, 5]

附：LeetCode 第295题参考代码：

解法1：

class MedianFinder:

    def __init__(self):
        """
        initialize your data structure here.
        """
        self.large_heap = []
        self.small_heap = []

    def addNum(self, num: int) -> None:
        if len(self.large_heap) == len(self.small_heap):
            heappush(self.small_heap, num)
            heappush(self.large_heap, - heappop(self.small_heap))
        else:
            heappush(self.large_heap, - num)
            heappush(self.small_heap, - heappop(self.large_heap))

    def findMedian(self) -> float:
        if len(self.large_heap) == len(self.small_heap):
            return (- self.large_heap[0] + self.small_heap[0]) / 2
        else:
            return - self.large_heap[0]


if __name__ == '__main__':
    a = MedianFinder()
    a.addNum(-1)
    print(a.findMedian())

    a.addNum(-2)
    print(a.findMedian())

    a.addNum(-3)
    print(a.findMedian())

    a.addNum(-4)
    print(a.findMedian())

    a.addNum(-5)
    print(a.findMedian())

解法二（使用内置库）

#!/usr/bin/python

import unittest

import heapq


class MedianFinder:

    def __init__(self):
        """
        initialize your data structure here.
        为了方便直接使用内置堆队列库heapq
        否则需自己构建通过数组建对二叉堆(以大根堆为例)：
        从最后的非叶子节点开始递归，若父亲节点<子节点中最大值，则将其交换。直到符合大根堆的定义(堆顶即数组下标为0的值要最大)
        """
        self.large_heap = []
        self.small_heap = []

    def add_num(self, num: int) -> None:
        if len(self.large_heap) == len(self.small_heap):
            heapq.heappush(self.small_heap, num)
            heapq.heappush(self.large_heap, - heapq.heappop(self.small_heap))
        else:
            heapq.heappush(self.large_heap, - num)
            heapq.heappush(self.small_heap, - heapq.heappop(self.large_heap))

    def find_median(self) -> float:
        if len(self.large_heap) == len(self.small_heap):
            return (- self.large_heap[0] + self.small_heap[0]) / 2
        else:
            return - self.large_heap[0]


class TestMedian(unittest.TestCase):
    def testMedian(self):
        m = MedianFinder()
        m.add_num(1)
        m.add_num(2)
        m.add_num(3)
        result = m.find_median()
        self.assertEqual(2, result, 'expect median is ：1.5')
        print('\n' * 2)
        print("median result is %s" % result)


if __name__ == '__main__':
    unittest.main()

使用heapq系统库执行用时更短…可能是因为heapq里调了c实现…

总结

• 函数里，见名知意的命名能够更好的理清编码时的逻辑