Python3-[39]面试题40.最小的k个数（TOPK问题）A

题目

输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。
P.S. 题目来源于leetcode

解题思路

这是TopK的问题，参看TOPK讲解。
具体来说，为求解最小的k个数，可以使用2中方法：大顶堆和快速排序。下面分别利用大顶堆和快速排序得到TopK.

解法1：heap

由于在Python3中提供的heapq是小顶堆，我们可以使用这样的策略得到大顶堆:
假设我们有数组lists

• 列表转heap:

lists = [-x for x in lists]
heapq.heapify(lists)

• 压入元素x：

heapq.heappush(lists, -x)

• 弹出元素并使用或查找最大元素：

y = - heapq.heappop(lists)#弹出并使用堆元素
maxNum = -lists[0]#找到最大元素

使用大顶堆找到最小的K个元素的思路：
建立size=k的最大堆，遍历其余的元素，如果有元素小于堆中的最大元素，则替代堆中的最大元素.如，堆中的元素为[6,4,3,2],其余元素为[5]，则有得到的最小的k个元素为[5,4,3,2]

code1：max-heap

import heapq
def TopK(nums,k):
	if len(nums)<=k: return nums
	heap = [-x for x in nums[:k]]
	heapq.heapify(heap)
	for num in nums[k:]:
		if num < -heap[0]:
			heapq.heappop(heap)
			heapq.heappush(heap,-num)
	return [-x for x in heap]

解法2：快速排序

我们可以使用快排的思路，只不过我们要修改快排的基准值base_idx = 0

• 递归基：
  • k==0时应返[]；
  • len(arr)==k时，应返回arr
• 分割：
  • base->基准arr[0]；
  • la:小于等于base的数组；
  • ra:大于base的数组
• 尾递归：
  • 当la的长度不够k时，base和ra来补充，ra补充前k-len(la)-1个元素，即partition(ra,k-len(la)-1)
  • 当la的长度超过k时，对la取前k个元素，即partition(la,k)

code2:快排思路

def TopK(nums,k):
	if k<1: return []
	if len(nums)<=k: return nums
	base = nums[0]
	la = [num for num in nums[1:] if num<=base]
	ra = [num for num in nums[1:] if num>base]
	if len(la)<k: return la + [base] + TopK(nums,k-len(la)-1)
	else: return TopK(la,k)