堆排序-topk问题

@小森

发布于 2024-03-15 10:25:35

780

发布于 2024-03-15 10:25:35

文章被收录于专栏：xiaosenxiaosen

现在有n个数，设计算法得到前k大的数

解决思路：排序后切片 O(nlogn)

但如果有上万个元素，只取前几个，就造成大量浪费

如果使用冒泡排序，则需要只执行前k趟冒泡（选择排序，插入排序） O(kn)

更方便的方法就是使用堆排序: O(nlogk)

>取列表前k个元素建立一个小根堆。堆顶就是第k大的数。

>依次向后遍历原列表，对于列表中的元素，如果小于堆顶，则忽略该元素；如果大于堆顶，则将堆顶更换为该元素，并且对堆进行一次调整。

>遍历列表所有元素后，倒序弹出堆顶。

调整部分的代码

跟堆排序的代码一样，只是两处符号变反

def sift(li,low,high):
    i=low
    j=2*i+1
    tmp=li[low]
    while j<=high:
        if j+1<=high and li[j+1]<li[j]:
            j=j+1    # j指向右孩子
        if li[j]<tmp:
            li[i]=li[j]
            i=j
            j=2*i+1
        else:
            li[i]=tmp
            break
    else:
        li[i]=tmp

其余代码:

def topk(li,k):
    heap=li[0:k]
    for i in range((k-2)//2,-1,-1):
        sift(heap,i,k-1)
        # 以上为建堆
    for i in range(k,len(li)-1):
        if li[i]>heap[0]:
            heap[0]=li[i]
            sift(heap,0,k-1)
            # 以上为遍历
    # 开始把小根堆遍历出来
    for i in range(k-1,-1,-1):
        # i一直指向最后一个元素
        heap[0], heap[i] = heap[i], heap[0]
        sift(heap,0, i - 1)
    return heap

本文参与腾讯云自媒体同步曝光计划，分享自作者个人站点/博客。

原始发表：2023-12-28，如有侵权请联系 cloudcommunity@tencent.com 删除

heap