使用快速排序、堆和桶解决「TopK问题」（Python）_topk 桶排序

TopK问题即求解第K个最大or最小元素的问题，一般使用快排、堆和桶来实现。

1 快速排序：

先总结快速排序的python模板。首先使用partition函数完成对数组从pivot的划分工作。

partition函数模板：

# 将数组分为小于基准值、基准值、大于基准值三个部分
def partition(nums, left, right):
    pivot = nums[left]  # 初始化一个基准值
    i, j = left, right
    while(i < j):
        while(i < j and nums[j] >= pivot): j -= 1  # 从后往前查找，直到找到一个比pivot小的数
        nums[i] = nums[j]  # 把这个数放左边来
        while(i < j and nums[i] <= pivot): i += 1  # 从前往后查找，直到找到一个比pivot大的数
        nums[j] = nums[i]  # 把这个数放右边来
    nums[i] = pivot 
    return i  # 返回基准值最终位置
'
运行
运行

然后使用partition递归完成快排。

快排模板：

def quicksort(nums, left, right):
    if left < right:
        index = partition(nums, left, right)  # 一个一个固定pivot的位置，直到每个数字位置都固定好
        quicksort(nums, left, index - 1)
        quicksort(nums, index + 1, right)
'
运行
运行

而topk问题则就是快排并不需要全部完成，只完成到第k个位置的数被固定的时候就可以停止了。

top-k切分：

将快速排序改成快速选择，也就是找到应该固定在这个数组排好序之后的第k个位置上的那个数。这个数左边是比它更小的数，右边是比它大的数。找到该位置后停止迭代，完成一次划分。

def topk_split(nums, k, left, right):
    if left < right:
        index = partition(nums, left, right)
        if index == k: return
        elif index < k: topk_split(nums, k, index+1, right)
        else: topk_split(nums, k, left, index - 1)
'
运行
运行

因此，可以根据topk_split解决一系列top-k问题：

获取前k小的数：

def topk_small(nums, k, left, right):
    topk_split(nums, k, 0, len(nums) - 1)
    return nums[:k]
'
运行
运行

获取前k大的数：

def topk_big(nums, k, left, right):
    topk_split(nums, len(nums) - k, 0, len(nums) - 1)  # 从len(nums)-k处划分好后右边部分就是前k大的数
    return nums[len(nums) - k:]
'
运行
运行

只排序前k个小的数：

def topk_sort_left(nums, k, left, right):
    topk_split(nums, k, 0, len(nums) - 1)  # 从k划分，取前k个数快排
    topk = nums[:k]
    quicksort(topk, 0, k - 1)
    return topk + nums[k:]
'
运行
运行

只排序后k个大的数：

def topk_sort_right(nums, k, left, right):
    topk_split(nums, k, 0, len(nums) - 1)
    topk = nums[k:]
    quicksort(topk, k, len(nums) - 1)
    return nums[:k] + topk
'
运行
运行

（注意python列表切片包前不包后，[start: stop: step]中，包括start但不包括stop）

2 堆

用于求解top-k问题。例如，求第k个最大元素，则建立小顶堆，不断往小顶堆中push新元素，当堆中元素数量>k时，pop堆顶元素，最后剩下的即为最大的k个元素，堆顶为第k个最大元素（求第k个最小元素就建立大顶堆）。

Python heapq库：

提供小顶堆的构建和一些对小顶堆的基本操作方法(如入堆，出堆等)，可以用于实现堆排序算法（大顶堆可以取相反数来实现）。

创建堆的方法：

import heapq

1. heappush(heap, num)，先创建一个空堆，然后将数据一个一个地添加到堆中。每添加一个数据后，heap都满足小顶堆的特性。

2. heapify(array)，直接将数据列表调整成一个小顶堆。

常用方法：

heappop(heap), heappush(heap, num), nlargest(n, heap), nsmallest(n, heap), merge(list1, list2)

使用heapq实现堆排序：

array = [10, 17, 50, 7, 30, 24, 27, 45, 15, 5, 36, 21]
# 使用方法1创建堆
heap = []
for num in array:
    heappush(heap, num)
heap_sort = [heappop(heap) for i in range(len(heap))]
print("heap sort result: ", heap_sort)

获取堆中top-k最大或最小值：

array = [10, 17, 50, 7, 30, 24, 27, 45, 15, 5, 36, 21]
# 使用方法2创建堆
heapify(array)
print(nlargest(2, array))
print(nsmallest(3, array))

归并两个有序列表：

array_a = [10, 7, 15, 8]
array_b = [17, 3, 8, 20, 13]
array_merge = merge(sorted(array_a), sorted(array_b))  # 返回迭代器
print("merge result:", list(array_merge))

其他方法：

heappushpop(heap, num)：先push num，再pop堆顶元素。

heapreplace(heap, num)：先pop堆顶元素，再push num。

3 桶

使用桶解决出现频率最多的k个元素问题。设置若干个桶，每个桶存储出现频率相同的数。桶的下标表示数出现的频率，即第i个桶中存储的数出现的频率为i。当把数都放入桶后，从后向前遍历桶，最先得到的 k 个数就是出现频率最多的的 k 个数。

一般需要先使用dict记录数字出现的频率，再根据频率放入对应频率的桶里。如果可以有相同频率的数，那么桶就是一个二维数组。

# 题目：给定数组nums和k，输出nums中出现频率前k高的元素。（力扣347）
class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        freq = {}  # dict记录数字出现频率
        for i in nums:
            if i in freq: freq[i] += 1
            else: freq[i] = 1
        bucket = [None] * (len(nums) + 1)  # 桶序号代表频率，每个桶保存相同频率的数组成的列表
        for i in freq.keys():
            if not bucket[freq[i]]: 
                bucket[freq[i]] = []
            bucket[freq[i]].append(i)
        ans = []
        for i in range(len(bucket) - 1, -1, -1):  # 从后往前输出k个数
            if not bucket[i]: continue
            ans.extend(bucket[i])
            if len(ans) == k: return ans
        return ans

参考：

Leetcode 题解 - 排序 | CS-Notes

力扣

、Python heapq库的用法介绍_小斌哥ge的博客-CSDN博客_heapq pytho