如何改进此代码以查找数组的 k 个最大元素？

Question

以下用于查找数组的 k 个最大元素的代码导致 TLE 错误。我如何优化它以使其运行更快？

import heapq    
for _ in range(int(input())):         
    n,k=map(int,input().split())   
    lists=list(map(int,input().split()))  
    
    heapq.heapify(lists)       
    
    for i in range(k+1):
        klargest=heapq.nlargest(i,lists)  
    
    print(*klargest)

Answer 1

for i in range(k+1):
   klargest=heapq.nlargest(i,lists)

每个 klargest 操作的时间复杂度是 O(k*log n)) 其中 n 是堆中元素的数量。在上面的代码片段中，对于值 [0,k].

，此操作是运行 k+1 次

正在计算循环时间：

迭代值（时间）

i == 0 (0*log(n))

i == 1 (1*log(n))

i == 2 (2*log(n))

.....

i == k-1 ((k-1)*log(n))

i == k ((k)*log(n))

总时间将是每个操作所用时间的总和 = (0.log(n)) + (1*log(n)) + .... + ((k-1)* log(n)) + ((k)*log(n))

总时间 = (0+1+2...+(k-1)+k)log(n) = ((k(k+1))/2)*log(n)

Total time ~~ O(k^2*(log(n)))

That's why the above code results in TLE.

优化方法：

import heapq    
for _ in range(int(input())):         
    n,k=map(int,input().split())   
    lists=list(map(int,input().split()))  
    
    heapq.heapify(lists)       
    
    for i in range(n-k):
        heapq.heappop(lists)
    klargest = list(lists) # converting heap to list
    print(*klargest)

因为 python 中的内置堆是最小堆。所以上面的代码从列表中弹出最少的 n-k 个元素。弹出每个操作将花费 log(n) 时间。因此总时间将是 ~~ (n-k)*logn。堆中剩下的k个元素就是我们要找的k大元素

因此，上述解决方案的时间复杂度为 O((n-k)*log(n)) == O(nlog(n)) 这是优化时间复杂度。

如何改进此代码以查找数组的 k 个最大元素？

How to improve this code for finding the k largest elements of an array?

algorithm

performance

priority-queue

binary-heap

heapq