堆排序 - topk 问题 (n 特别大时,很快)
现在有 n 个数,设计算法得到前 k 大的数
解决思路:
1. 取列表前 k 个元素建立一个小根堆,堆顶就是目前第 k 大的数
2. 依次向后遍历原列表,对于列表中的元素,如果小于堆顶,则忽略该元素
如果大于堆顶,则将堆顶更换为该元素,并对堆进行一次调整
3. 遍历列表所有元素后,倒序弹出堆顶
# 比较排序 - 通过比较来决定两个值是否交换
import random
def sift(li, low, high):
i = low
j = 2 * i + 1
temp = li[low]
while j <= high:
if j + 1 <= high and li[j + 1] < li[j]:
j = j + 1
if li[j] < temp:
li[i] = li[j]
i = j
j = 2 * i + 1
else:
break
li[i] = temp
def topk(li, k):
# 提取前 k 个元素建堆
heap = li[0:k]
for i in range((k - 2) // 2, -1, -1):
sift(heap, i, k - 1)
# 遍历剩下的元素,如果比堆顶大,就加进去,并对堆进行一次调整
for i in range(k, len(li) - 1):
if li[i] > heap[0]:
heap[0] = li[i]
sift(heap, 0, k - 1)
# 挨个出数
for i in range(k - 1, -1, -1):
heap[0], heap[i] = heap[i], heap[0]
sift(heap, 0, i - 1)
return heap