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)),这是优化的时间复杂度。