从数组中找到topK的元素(序号)

问题:
在n个数中找出最大的k个数。

多次求min()或求max()

最简单的方法是对大小为k的数组进行n次求min计算(或者对大小为n的数组进行k次求max计算)最后能够找出最大k个数。复杂度是O(nk)。
代码:

def topK_mink(num_list,k):
    topK=num_list[:k]
    for i in range(k,len(num_list)):
        topK_min=min(topK)
        if num_list[i]>topK_min:
            topK[topK.index(topK_min)]=num_list[i]
    return topK        

使用小根堆

维护一个大小为k的小根堆,从头到尾扫描n个数,如果当前数比堆顶大,替换堆顶,这样扫描到最后堆中保存的是最大的k个数。复杂度是O(nlogk)
代码:

import heapq
def topK_heapq(num_list,k):
    array = []
    for i in range(len(num_list)):
        if len(array) < k:
            heapq.heappush(array, num_list[i])
        else:
            array_min = array[0]
            if num_list[i] > array_min:
                heapq.heapreplace(array, num_list[i])
    topK=array
    return topK

使用大根堆

维护一个大小为n的大根堆,每次弹出堆顶元素,共弹出k次。复杂度O(klogn)
代码:略

快速选择BFPRT

借用快速排序中思想,在快排中每次用一个轴将数组划分为左右两部分,轴左边的数都小于轴,轴右边的数都大于轴,轴所在的位置和排好序后的位置相同。这里只要找到第k大的数作为轴进行划分,那么就找到了最大的k个数。期望复杂度是:O(n)
代码:

def topK_partition(arr,k):
    
    def partition(num_list,left,right,k):
        flag=num_list[left]
        i=left
        j=right
        while i<j:
            #print(flag,i,j,num_list)
            if num_list[i]>flag:
                i+=1
            elif num_list[j]<flag:
                j-=1
            else:
                if num_list[i]==num_list[j]:
                    j-=1
                num_list[i],num_list[j]=num_list[j],num_list[i]
        #print(flag,num_list)
        if i<k:
            return partition(num_list,i+1,right,k)
        if i>k:
            return partition(num_list,left,i-1,k)
        return num_list[:k]
    
    return partition(arr[:],0,len(random_list)-1,k)

测试代码:

import numpy as np
import time
def judge(ans,k1):
    for i in ans:
        if i-k1<0:
            return False
    return True
k=1000
n=25000
random_list=[np.random.randint(n*0.5) for i in range(n)]
real_ans=sorted(random_list,reverse=True)[:k+1]
k1=real_ans[-1]

t1=time.time()
ans=topK_heapq(random_list,k)
t2=time.time()
print(judge(ans,k1),t2-t1)

t1=time.time()
ans=topK_partition(random_list,k)
t2=time.time()
print(judge(ans,k1),t2-t1)

t1=time.time()
ans=topK_mink(random_list,k)
t2=time.time()
print(judge(ans,k1),t2-t1)

结果:
topK_mink()没有任何优势
topK_partition()的运行时间不稳定
topK_heapq()运行时间稳定
虽然期望复杂度topK_partition()优于topK_heapq(),但是topK_partition()计算开销比topK_heapq()多。
当n小时,用topK_heapq()比topK_partition()好
当n大,k小时,topK_heapq()用时也较短。
当n大,k大时(n>10,000,000),用topK_partition()。

posted @ 2017-03-02 17:01  机器狗mo  阅读(4703)  评论(0编辑  收藏  举报