python heapq库原理及使用

一、heapq堆的原理

1.1、堆的本质

1. 它是一个完全二叉树
2. 实现的时候我们不需要建造一个树，改用一个数组即可
   给树的节点编号，代表着数组中元素的下边。
   
   由此我们可以得出结论：
3. 已知节点的编号为index，他的父节点的编号：father_index = (index - 1) / 2
4. 左孩子节点的编号为：left_index = index * 2 + 1
5. 右孩子节点的编号为：right_index = index * 2 + 2

1.2、如何调整堆

1. 添加元素

• 把新数据添加到树的最后一个元素，也就是数组的末尾
• 把末尾节点向上调整

2. 弹出堆顶

• 交换根节点与最后一个节点的值
• 删除最后一个节点
• 把根节点向下调整

1.3、堆的特点

1. 内部数据是有序的
2. 可以弹出堆顶的元素，大顶堆就是弹出最大值，小顶堆就是弹出最小值
3. 每次加入新元素或者弹出堆顶元素后，调整堆使之重新有序仅需要O(logn)的时间
4. 支持在线算法

二、堆的基本使用

import heapq 
#heapq是默认构造小顶堆
nums = [4, 3, 2, 1]
hp = []
for num in nums:
	#构建小顶堆hp，并插入元素
    heapq.heappush(hp, num)
print(hp)  # [1, 2, 3, 4]

#直接构建小顶堆,并转换成列表
nums = [4, 3, 2, 1]
heapq.heapify(nums)
print(nums) # [1, 3, 2, 4]

#弹出堆顶元素，重新调整堆
nums = [4, 3, 2, 1]
heapq.heapify(nums)
heapq.heappop(nums)
print(nums) # [2, 3, 4]

#弹出最小元素，并插入元素调整堆的位置
nums = [4, 3, 2, 1]
heapq.heapify(nums)
heapq.heapreplace(nums, 5)
print(nums) # [2, 3, 5, 4]

#返回出前K个最大元素
nums = [4, 3, 2, 1]
heapq.heapify(nums)
k = heapq.nlargest(2, nums)
print(k) # [4, 3]

#返回前K个最小元素
nums = [4, 3, 2, 1]
heapq.heapify(nums)
k = heapq.nsmallest(2, nums)
print(k) # [1, 2]

三、简单的实现堆

class Heap:
    def __init__(self,desc=False):
        """
        初始化，默认创建一个小顶堆
        """
        self.heap = []
        self.desc = desc
    
    @property
    def size(self):
        return len(self.heap)
    
    def top(self):
        if self.size:
            return self.heap[0]
        return None
    
    def push(self,item):
        """
        添加元素
        第一步，把元素加入到数组末尾
        第二步，把末尾元素向上调整
        """
        self.heap.append(item)
        self._sift_up(self.size-1)
    
    def pop(self):
        """
        弹出堆顶
        第一步，记录堆顶元素的值
        第二步，交换堆顶元素与末尾元素
        第三步，删除数组末尾元素
        第四步，新的堆顶元素向下调整
        第五步，返回答案
        """
        item = self.heap[0]
        self._swap(0,self.size-1)
        self.heap.pop()
        self._sift_down(0)
        return item
    
    def _smaller(self,lhs,rhs):
        return lhs > rhs if self.desc else lhs < rhs
    
    def _sift_up(self,index):
        """
        向上调整
        如果父节点和当前节点满足交换的关系
        （对于小顶堆是父节点元素更大，对于大顶堆是父节点更小），
        则持续将当前节点向上调整
        """
        while index:
            parent = (index-1) // 2
            
            if self._smaller(self.heap[parent],self.heap[index]):
                break
                
            self._swap(parent,index)
            index = parent
    
    def _sift_down(self,index):
        """
        向下调整
        如果子节点和当前节点满足交换的关系
        （对于小顶堆是子节点元素更小，对于大顶堆是子节点更大），
        则持续将当前节点向下调整
        """
        # 若存在子节点
        while index*2+1 < self.size:
            smallest = index
            left = index*2+1
            right = index*2+2
            
            if self._smaller(self.heap[left],self.heap[smallest]):
                smallest = left
                
            if right < self.size and self._smaller(self.heap[right],self.heap[smallest]):
                smallest = right
                
            if smallest == index:
                break
            
            self._swap(index,smallest)
            index = smallest
    
    def _swap(self,i,j):
        self.heap[i],self.heap[j] = self.heap[j],self.heap[i]```