从数组中找到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()。