时间戳

首页 > 解决方案 > 如何改进此代码以查找数组的 k 个最大元素？

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

问题描述

以下用于查找数组的 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)

标签： algorithmperformancepriority-queuebinary-heapheapq

解决方案

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)

总时间~~ O(k^2*(log(n)))

这就是为什么上面的代码会导致 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中的内置堆是最小堆。所以上面的代码是从列表中弹出最少 nk 个元素。弹出每个操作将花费 log(n) 时间。因此总时间将为 ~~ (nk)*logn。堆中剩余的 k 个元素是我们想要找到的 k 个最大元素。

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

推荐阅读