Python高级操作第一弹：heapq的妙用

关注

原创最新推荐文章于 2026-03-15 00:56:39 发布 659 阅读 AI 写同款· GEO 优化 ›

1 基础操作

"""
1）基础：入堆、出堆、替换操作；heapq.heappush(heapq, item)/heapq.heappop(heapq, item)/
        .heappushpop(压入一个，弹出一个）/.heapreplace（先弹出堆顶，再压入元素，和heappushpop顺序相反）
   初始化：创建堆用[]的列表，或者heapify()来将填充的列表转化为堆
   操作：.merge(迭代对象，key,reverse)合并多个输入，.nlarget(n,迭代，key=none） .nsmallest（n) n小时候有效，n大时候用sorted（）更好
2）heap 堆； 二叉堆 binary heap； Fibonacci heap； binomial heap 二项堆
3）更高级的：优先级堆（优先级队列实现方法）；——一种二叉树；heap[k] <= heap[2k+1], heap[k] <= heap[2k+2]
     heap[0]始终是全堆最小的值（如果压入-element那么就反过来最大的值，百试不爽）
"""
from heapq import *
from typing import List

import numpy as np
import copy

# ex1：
hq = []
heappush(hq, (2,'play game'))
heappush(hq, (1,'play with son'))
heappush(hq, (3,'invite his mom'))
heappush(hq, (4,'go out for lunch'))
[print(heappop(hq)) for i in range(4)]

# ex2:
hq_n = [4,3,5,6,9,100,-100,-200]
heapify(hq_n)
[print(heappop(hq_n)) for i in range(len(hq_n))]

print("inverse sequence")
hq_inv = copy.deepcopy([4,3,5,6,9,100,-100,-200])
hq_inv = list(-np.array(hq_inv))
heapify(hq_inv)
[print(-heappop(hq_inv)) for i in range(len(hq_inv))]
"""
Output:

(1, 'play with son')
(2, 'play game')
(3, 'invite his mom')
(4, 'go out for lunch')
-200
-100
3
4
5
6
9
100
inverse sequence
100
9
6
5
4
3
-100
-200
"""

这个玩意儿，简单一句话，就是用二叉树的方法实现了堆的树形结构，取用、排序稳定在O(nlogn)复杂度，最恶劣情况和平均情况相差无几；
然而，对于很多头疼的情况，比如实际中的磁盘存储平衡等，解决得很是巧妙；
这里碰到一个Leetcode第4题的期刊，对比下实现，喷了一口老血

2 LeetCode 004测试威力

2-1 title

“”"
LC04: 4. 寻找两个正序数组的中位数
给定两个大小为 m 和 n 的正序（从小到大）数组 nums1 和 nums2。请你找出并返回这两个正序数组的中位数。

进阶：你能设计一个时间复杂度为 O(log (m+n)) 的算法解决此问题吗？

示例 1：

输入：nums1 = [1,3], nums2 = [2]
输出：2.00000
解释：合并数组 = [1,2,3] ，中位数 2
示例 2：

输入：nums1 = [1,2], nums2 = [3,4]
输出：2.50000
解释：合并数组 = [1,2,3,4] ，中位数 (2 + 3) / 2 = 2.5
示例 3：

输入：nums1 = [0,0], nums2 = [0,0]
输出：0.00000
示例 4：

输入：nums1 = [], nums2 = [1]
输出：1.00000
示例 5：

输入：nums1 = [2], nums2 = []
输出：2.00000
“”"

2-2 传统的合并数组，删除合并掉的，思路直接操作繁琐25行代码

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> None:
        rst = []
        if len(nums1) + len(nums2) == 0:
            return None

        while nums1 and nums2:
            if nums1[0] < nums2[0]:
                rst.append(nums1[0])
                del nums1[0]
            else:
                rst.append(nums2[0])
                del nums2[0]

        while nums1:
            rst.append(nums1[0])
            del nums1[0]

        while nums2:
            rst.append(nums2[0])
            del nums2[0]

        x = len(rst)
        # print(rst, x//2)

        if x / 2 != x // 2:
            return rst[x // 2]
        else:
            return (rst[x // 2 - 1] + rst[x // 2]) / 2

        # return np.mean(rst[int((len(rst)/2-1)) : int(len(rst)/2)+1])
        # (float( rst[int((len(rst)/2-1))]) + float(rst[int(len(rst)/2)+1]) ) /2
        #很直接的思路，逐个加入元素，删掉原有的，比较耗费；2）注意结尾的中位数取值和长度的奇偶有关

2-3 扣奇偶数，算序号，烧脑，40+行

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        # 不删除元素，只存取i j序号；取到总长度len/2的元素和前一个元素
        rst = None
        length = len(nums1) + len(nums2)
        if length == 0:
            return None

        i, j = len(nums1) - 1, len(nums2) - 1

        if length // 2 != length / 2:
            seq = length // 2 + 1  # 找到第seq + 1个元素
            while i >= 0 and j >= 0:
                if nums1[i] > nums2[j]:
                    rst = nums1[i]
                    i -= 1
                else:
                    rst = nums2[j]
                    j -= 1
                seq -= 1
                if seq == 0:
                    return rst
            while i >= 0:
                i -= 1
                seq -= 1
                if seq == 0:
                    return nums1[i + 1]
            while j >= 0:
                j -= 1
                seq -= 1
                if seq == 0:
                    return nums2[j + 1]
        else:
            seq = length // 2 + 1  # 找第seq，和seq-1
            rst = [0, 0]
            while i >= 0 and j >= 0:
                seq -= 1
                if nums1[i] > nums2[j]:
                    rst[seq % 2] = nums1[i]
                    i -= 1
                else:
                    rst[seq % 2] = nums2[j]
                    j -= 1

                if seq == 0:
                    return (rst[0] + rst[1]) / 2

            while i >= 0:
                seq -= 1
                rst[seq % 2] = nums1[i]
                i -= 1
                if seq == 0:
                    return (rst[0] + rst[1]) / 2

            while j >= 0:
                seq -= 1
                rst[seq % 2] = nums2[j]
                j -= 1
                if seq == 0:
                    return (rst[0] + rst[1]) / 2
            # 改良的思路，开销比上一个还大；

    # python高级知识第一弹：heapq操作，堆操作；简化存储，快速构建数据结构的有效操作！

2-4 堆方法，10行，两个list加起来heapify一下，输出前一半，然后取末尾1-2个做中位数

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        # 建立两个堆，分别存放小值和大值；保持两个堆size一致（最多相差1）；
        #总数为偶数就找
        num = nums1 + nums2
        heapify(num)
        x = len(nums1) + len(nums2)
        if x == 0:
            return None

        L = nsmallest(x // 2 + 1, num)
        if x % 2: #奇数,x//2 + 1
            return max(L)
        else:
            return (L[-1] + L[-2])/2