数据结构-使用Python实现自己的堆

原创于 2025-05-27 14:06:16 发布 · 676 阅读

CC 4.0 BY-SA版权

文章标签：

Python自带heapq，默认是小顶堆，可以直接使用它快速完成想要的功能，比如力扣215. 数组中的第K个最大元素：

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        min_heap = []
        for num in nums:
            heapq.heappush(min_heap, num)
            if len(min_heap) > k:
                heapq.heappop(min_heap)
        return min_heap[0]

但是想要理解堆或者想要做一个大顶堆，还是要亲自动手维护一个堆才好。当热，大顶堆可以通过负数的小顶堆来实现。接下来就用Python手搓一个小顶堆。

首先我们应该明确堆应该具有哪些功能：1.插入新元素；2.删除并返回堆顶元素；3.返回堆顶元素但不删除。

1.插入新元素

class MinHeap:
    def __init__(self):
        self.heap = []

    def push(self, val):
        """插入新元素"""
        self.heap.append(val)
        self._sift_up(len(self.heap) - 1)
    
    def _sift_up(self, idx):
        """从 idx 向上调整"""
        parent = (idx - 1) // 2
        while idx > 0 and self.heap[idx] < self.heap[parent]:
            self.heap[idx], self.heap[parent] = self.heap[parent], self.heap[idx]
            idx = parent
            parent = (idx - 1) // 2

接下来我们写的函数都是在MinHeap这个类中，可以看到无论是heap还是stack，都离不开list。

这里牵扯到了_sift_up()这个函数，不是shift，而是sift，他的意思是向上筛。

这是由于我们向堆的末尾中添加了一个新的元素，需要根据堆的性质维护堆，即如果它比根节点小，就浮上去。

2.删除并返回堆顶元素

class MinHeap:
    ……

    def pop(self):
        """删除并返回堆顶元素（最小值）"""
        if not self.heap:
            raise IndexError("pop from empty heap")
        min_val = self.heap[0]
        last_val = self.heap.pop()
        if self.heap:
            self.heap[0] = last_val
            self._sift_down(0)
        return min_val

    
    def _sift_down(self, idx):
        """从 idx 向下调整"""
        n = len(self.heap)
        while True:
            smallest = idx
            left = 2 * idx + 1
            right = 2 * idx + 2

            if left < n and self.heap[left] < self.heap[smallest]:
                smallest = left
            if right < n and self.heap[right] < self.heap[smallest]:
                smallest = right

            if smallest == idx:
                break

            self.heap[idx], self.heap[smallest] = self.heap[smallest], self.heap[idx]
            idx = smallest

堆顶就是heap[0]，我们使用和维护堆的目的就是用它的堆顶。

这里牵扯到了_sift_down()这个函数，不是shift，而是sift，他的意思是向下筛。

这个函数的出现与堆的性质有关，当我们将堆顶取走后，我们需要把堆底最后一个元素拿到堆顶，然后再根据堆的性质维护一下堆，即如果它比子节点大，就沉下去。

3.返回堆顶元素但不删除

删除会做，不删除更简单，直接看一眼堆顶不对堆做任何操作即可：

class MinHeap:
    ……

    def peek(self):
        """返回堆顶元素但不删除"""
        if not self.heap:
            raise IndexError("peek from empty heap")
        return self.heap[0]

4.带有heapify的完整代码

不仅可以一个一个的插入，我们还可以获取一个list中的数，把这个list变成一个heap：

class MinHeap:
    def __init__(self, nums=None):
        self.heap = nums[:] if nums else []
        if self.heap:
            self._heapify()

    def _heapify(self):
        """将当前数组堆化成小顶堆"""
        n = len(self.heap)
        # 从最后一个非叶子节点开始向下调整
        for i in reversed(range(n // 2)):
            self._sift_down(i)

    def _sift_down(self, idx):
        n = len(self.heap)
        while True:
            smallest = idx
            left = 2 * idx + 1
            right = 2 * idx + 2
            if left < n and self.heap[left] < self.heap[smallest]:
                smallest = left
            if right < n and self.heap[right] < self.heap[smallest]:
                smallest = right
            if smallest == idx:
                break
            self.heap[idx], self.heap[smallest] = self.heap[smallest], self.heap[idx]
            idx = smallest

    def push(self, val):
        self.heap.append(val)
        self._sift_up(len(self.heap) - 1)

    def _sift_up(self, idx):
        parent = (idx - 1) // 2
        while idx > 0 and self.heap[idx] < self.heap[parent]:
            self.heap[idx], self.heap[parent] = self.heap[parent], self.heap[idx]
            idx = parent
            parent = (idx - 1) // 2

    def pop(self):
        if not self.heap:
            raise IndexError("pop from empty heap")
        min_val = self.heap[0]
        last = self.heap.pop()
        if self.heap:
            self.heap[0] = last
            self._sift_down(0)
        return min_val

    def peek(self):
        if not self.heap:
            raise IndexError("peek from empty heap")
        return self.heap[0]

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