Python堆的基本用法

官方链接：https://docs.python.org/3/library/heapq.html

堆，其实是二叉树的一种。Python中的堆其实是小根堆，即每一个父节点都小于等于它的子节点。

堆这种结构在leetcode刷题过程中经常遇到。

对于一个堆heap来说，它最小的元素总是它的根节点，即heap[0]。

对于一个堆来说，它的导入比较简单

import heapq

heap = []

heappush

通过heappush方法可以将数字加入堆中

heapq.heappush(heap,1)

heapq.heappush(heap,2)

heapq.heappush(heap,3)

heap
#[1,2,3]

heappop

通过heappop的方式，每次返回堆中的最小值。直到heap为空，报错

heapq.heappop(heap)
#1
heapq.heappop(heap)
#2
heapq.heappop(heap)
#3
heapq.heappop(heap)
Traceback (most recent call last):

  File "", line 1, in <module>
    heapq.heappop(heap)

IndexError: index out of range

heapify

另外，heapq还有heapify方法，把一个list转化为堆heap。

res = [1,4,3,5,6,2]

heapq.heapify(res)

res
#返回的res是二叉树结构的列表[1, 4, 2, 5, 6, 3]

nlargest

返回res中最大的两个数

res = [1,4,3,5,6,2]

heapq.nlargest(2,res)
#[6, 5]

nsmallest

返回res中最小的两个数

res = [1,4,3,5,6,2]

heapq.nsmallest(2,res)
#[1,2]

另外，heapq还存在merge（融合多个列表）、heapreplace、heappushpop等其他方法，具体可以查看文档了解。

堆在leetcode中的应用

41. 数据流中的中位数

• 思路：看到这道题，一个比较直观的想法就是，将所有的数字进行排序然后返回中位数。
  但是这道题只需要计算中位数，每次都排序的话复杂度会很高。
• 数据逻辑：引入堆结构，并构建两个堆来存储前半段的最大值（大根堆）和后半段的最小值（小根堆）。这样每次计算
  只要根据数据流的奇偶性对两个堆的根节点进行计算即可。

class MedianFinder(object):
    from heapq import *
    def __init__(self):
        """
        initialize your data structure here.
        """
        self.bigHpeap = []
        self.smallHeap = []
    def addNum(self, num):
        """
        :type num: int
        :rtype: None
        """
        if len(self.bigHpeap) == len(self.smallHeap):#总数为偶数时,先插入到大根堆,在插入到小根堆
            heapq.heappush(self.smallHeap, -heapq.heappushpop(self.bigHpeap, -num))
        else:#总数为奇数时,先插入到小根堆,在插入到大根堆
            heapq.heappush(self.bigHpeap, -heapq.heappushpop(self.smallHeap, num))

    def findMedian(self):
        """
        :rtype: float
        """
        if len(self.bigHpeap) == len(self.smallHeap):
            return (-self.bigHpeap[0] + self.smallHeap[0]) / 2.0
        else:
            return self.smallHeap[0]

23. 合并K个排序链表
合并K个排序链表，一种比较直接、暴力的方式就是把所有链表值append到一个列表中，排序后重新输出链表。

另一种方式是把每个链表的值放入heap，然后通过pop()每次自动排出最小值。

class Solution:
    def mergeKLists(self, lists: List[ListNode]) -> ListNode:
        import heapq
        dummy = ListNode(0)
        p = dummy
        head = []
        for i in range(len(lists)):
            while lists[i] :
                heapq.heappush(head, (lists[i].val, i))
                lists[i] = lists[i].next
        while head:
            val, idx = heapq.heappop(head)
            p.next = ListNode(val)
            p = p.next
        return dummy.next