堆的实现与TopK问题

喜欢我的文章，欢迎关注个人公众号：数学编程

---

在《图文讲解堆排序》中介绍了堆数据的基本概念以及存储方法，同时也介绍了其应用就是堆排序。虽然堆的概念相对简单，代码的实现具有学习意义。本文重点介绍最小堆的实现，再介绍一种堆的应用，解决TopK问题。

Talk is cheap, Show me the code

很多时候原理简单，代码的实现并非如此，下面就重点解读一下代码。本文代码参考heapq.py源码

最小堆的实现

首先来看一下最小堆的抽象数据类型（Abstract Data Type， ADT），用代码呈现如下：

class MinHeap:
    def heapify(self): # 堆化
        pass
    def _sift_up(self, index): # 向上筛选
        pass
    def _sift_down(self, start_pos: int, pos: int): # 向下筛选
        pass
    def insert(self, item: int): # 插入元素
        pass
    def pop(self): # 弹出堆顶
        pass

上面列出了比较重要的5种操作方法，最核心的就是向上筛选和向下筛选两种，接下来重点介绍这两种操作：

向上筛选

向上筛选是将子节点中最小的数据向上筛选的过程（最大堆就向上筛选最大数）。等等，我们不是在讲最小堆吗？最小的数据不是在堆顶吗？对的，那么向上筛选的动作就是删除堆顶以后，将最后一个元素放在堆顶，然后执行向上筛选，找出最小的元素作为新的堆顶。

待筛选的元素通常是堆顶元素，目标是从子节点中找到最小的来作为新的顶堆。while循环中与子节点比较，找到待筛选元素item的合适位置，最后再做一次向小筛选。

def _sift_up(self, index):
   """
   向上筛选 与子节点对比
   :param index: 待筛选元素的索引
   """
   item = self._heap[index]
   pos = index
   child_pos = 2 * pos + 1
   while child_pos < len(self._heap):
       right_pos = child_pos + 1

       if right_pos < len(self._heap) and \
               self._heap[child_pos] >= self._heap[right_pos]:
           child_pos = right_pos  # 保证child_pos是最小的

       self._heap[pos] = self._heap[child_pos]
       pos = child_pos
       child_pos = 2 * pos + 1

   self._heap[pos] = item
   self._sift_down(index, pos)  # 向下筛选

向下筛选

向下筛选是与父节点比较，把较大的元素筛选下来。（对最大堆来说就是把较小的元素筛选下来）。在最小堆中父节点一定是小于子节点的值，那么向下筛选的使用场景就是添加新元素，然后不断与其父节点比较，直到找到大于父节点的位置放入。

    def _sift_down(self, start_pos: int, pos: int):
        """
        向下筛选跟父节点的值比较
        :param start_pos: 开始位置，通常为堆顶（0）
        :param pos: 筛选的终止位置
        """
        item = self._heap[pos]
        while start_pos < pos:
            parent_pos = (pos - 1) >> 1
            parent = self._heap[parent_pos]
            if parent > item:
                self._heap[pos] = parent
                pos = parent_pos
            else:
                break
        self._heap[pos] = item

TopK问题描述

排序在生活中很常见，把原本无序的排列变得有序。但有时候我们不需要对所有的元素进行排序，比如世界500强企业，十大杰出青年等等。我们只关心最多或者最少的部分元素，这个就是TopK问题。

要解决这样的TopK问题，其实思路有很多，最直接的解决办法就是从序列拿出最大（最小）的元素，然后重复操作K次，就解决了这个问题。最直接的方法往往并不是最有效的，如果问题的规模为N，那么解决这个问题的时间复杂度将是$o(NK)$。使用效果最好的排序算法，其时间复杂度将是$o(N\log N)$.

问题在于我们只关心前Top K个元素，并不关心所有元素的顺序，排序再取Top K计算量多余了。假如我们有N个元素，需要找到前K个大的元素。我们只需要定义存放K个元素长度的最小堆，然后遍历剩余的元素，与堆顶比较，如果大于对顶元素，则替换堆顶元素，然后堆化，重复比较剩余元素，直到结束，总结起来就是：

1. 创建K个元素的最小堆
2. 遍历剩余元素与堆顶比较，如果大于堆顶，则替换堆顶，执行向上筛选操作
3. 重复2步骤直到遍历结束
4. 输出堆中的K个元素

分析一下算法时间复杂度，获取堆顶元素的复杂度为$o(1)$，K个元素执行一次堆化操作时间复杂度为$o(\log K)$，如果对每个元素都执行堆化操作，那么算法的时间复杂度为$o(N\log K)$效果明显好很多。

总结

本文在图文讲解堆排序的文章基础上进一步介绍了堆数据的实现，其中核心的是向上筛选和向下筛选操作。本来文章会提前一天发出来，由于本人实现的代码与参考文献的源码存在差距，毕竟源码经过长时间的维护，各方面效果都更好。同时介绍了TopK问题，以及解决思路。

排序或者TopK问题，在大数据算法中非常常见，面对海量数据要找出TopK，已经超出了单机所能处理的范围，需要借助集群做分布式计算了，但是其基本的原理是一致的。