算法与数据结构（二十五）TopK问题：基于快排的Python模板

首先，先写partition模板

def partition(nums, left, right):pivot = nums[left]#初始化一个待比较数据i,j = left, rightwhile(i < j):while(i<j and nums[j]>=pivot): #从后往前查找，直到找到一个比pivot更小的数j-=1nums[i] = nums[j] #将更小的数放入左边while(i<j and nums[i]<=pivot): #从前往后找，直到找到一个比pivot更大的数i+=1nums[j] = nums[i] #将更大的数放入右边#循环结束，i与j相等nums[i] = pivot #待比较数据放入最终位置 return i #返回待比较数据最终位置

1. 快速排序

复习一下快速排序：

#快速排序
def quicksort(nums, left, right):if left < right:index = partition(nums, left, right)quicksort(nums, left, index-1)quicksort(nums, index+1, right)arr = [1,3,2,2,0]
quicksort(arr, 0, len(arr)-1)
print(arr)

2. topk切分

将快速排序改成快速选择，即我们希望寻找到一个位置，这个位置左边是k个比这个位置上的数更小的数，右边是n-k-1个比该位置上的数大的数，我将它命名为topk_split，找到这个位置后停止迭代，完成了一次划分。

def topk_split(nums, k, left, right):#寻找到第k个数停止递归，使得nums数组中index左边是前k个小的数，index右边是后面n-k个大的数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问题

3. 获得前k小的数

#获得前k小的数
def topk_smalls(nums, k):topk_split(nums, k, 0, len(nums)-1)return nums[:k]arr = [1,3,2,3,0,-19]
k = 2
print(topk_smalls(arr, k))
print(arr)

4. 获取第k小的数

#获得第k小的数
def topk_small(nums, k-1):topk_split(nums, k, 0, len(nums)-1)return nums[k] arr = [1,3,2,3,0,-19]
k = 3
print(topk_small(arr, k))
print(arr)

5. 获得前k大的数

#获得前k大的数 
def topk_larges(nums, k):#parttion是按从小到大划分的，如果让index左边为前n-k个小的数，则index右边为前k个大的数topk_split(nums, len(nums)-k, 0, len(nums)-1) #把k换成len(nums)-kreturn nums[len(nums)-k:] arr = [1,3,-2,3,0,-19]
k = 3
print(topk_larges(arr, k))
print(arr)

6. 获得第k大的数

#获得第k大的数 
def topk_large(nums, k):#parttion是按从小到大划分的，如果让index左边为前n-k个小的数，则index右边为前k个大的数topk_split(nums, len(nums)-k, 0, len(nums)-1) #把k换成len(nums)-kreturn nums[len(nums)-k] arr = [1,3,-2,3,0,-19]
k = 2
print(topk_large(arr, k))
print(arr)

7. 只排序前k个小的数

#只排序前k个小的数
#获得前k小的数O(n)，进行快排O(klogk)
def topk_sort_left(nums, k):topk_split(nums, k, 0, len(nums)-1) topk = nums[:k]quicksort(topk, 0, len(topk)-1)return topk+nums[k:] #只排序前k个数字arr = [0,0,1,3,4,5,0,7,6,7]
k = 4
topk_sort_left(arr, k)

8. 只排序后k个大的数

#只排序后k个大的数
#获得前n-k小的数O(n)，进行快排O(klogk)
def topk_sort_right(nums, k):topk_split(nums, len(nums)-k, 0, len(nums)-1) topk = nums[len(nums)-k:]quicksort(topk, 0, len(topk)-1)return nums[:len(nums)-k]+topk #只排序后k个数字arr = [0,0,1,3,4,5,0,-7,6,7]
k = 4
print(topk_sort_right(arr, k))