剑指Offer-40 最小的K个数（快排变形，堆）

题目描述

输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。

示例 1：

输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]

示例 2：

输入：arr = [0,1,2,1], k = 1
输出：[0]

限制：
    0 <= k <= arr.length <= 10000
    0 <= arr[i] <= 10000

题解一（快排思想）

对于数组的 [left, right] 区间，利用快排的思想，选择一个起始数，让比起小的在它左边，比其大的在右边。

使用的方法是填坑法（这位博主写的很好）：

int choose = arr[left];
int p1 = left, p2 = right;
while (p1 < p2)
{
	while (p1 < p2 && arr[p2] >= choose) p2--;
	if (p1 < p2) {
		arr[p1] = arr[p2];
	}
	while (p1 < p2 && arr[p1] <= choose) p1++;
	if (p1 < p2) {
		arr[p2] = arr[p1];
	}
}
arr[p1] = choose;

循环结束后，左指针和右指针指向了同一个位置，任意判断其一：

• 如果下标刚好等于k - 1，那么可以返回从 [0, k - 1] 一共k个数。

• 如果下标比k - 1大，那么说明前k个数在左边，递归进行[left, p - 1]处理。

• 如果下标比k - 1小，那么说明前k个数在右边，递归进行[p + 1, right]处理。

完整代码实现：

public int[] getLeastNumbers(int[] arr, int k) {
	if (k == 0 || arr == null || arr.length == 0) {
		return new int[]{};
	}

	quickSelect(0, arr.length - 1, arr, k);
	return Arrays.copyOf(arr, k);
}

private void quickSelect(int left, int right, int [] arr, int k)
{
	int choose = arr[left];
	int p1 = left, p2 = right;
	while (p1 < p2)
	{
		while (p1 < p2 && arr[p2] >= choose) p2--;
		if (p1 < p2) {
			arr[p1] = arr[p2];
		}
		while (p1 < p2 && arr[p1] <= choose) p1++;
		if (p1 < p2) {
			arr[p2] = arr[p1];
		}
	}
	arr[p1] = choose;

	if (p1 == k - 1) {
		return;
	} else if (p1 < k - 1) {
		quickSelect(p1 + 1, right, arr, k);
	} else {
		quickSelect(left, p1 - 1, arr, k);
	}
}

时间复杂度：$O(N)$。
假设要找下标为k的元素，第一次切分的时候需要遍历整个数组 (0 ~ n) 找到了下标是 j 的元素，假如 k 比 j 小的话，那么我们下次切分只要遍历数组 (0 ~ k-1)的元素就行啦，反之如果 k 比 j 大的话，那下次切分只要遍历数组 (k+1～n) 的元素就行啦，总之可以看作每次调用 partition 遍历的元素数目都是上一次遍历的 $1/2$，因此时间复杂度是 $N+N/2+N/4+...+N/N=2N$, 因此时间复杂度是 $O(N)$。

空间复杂度：$O(log(N))$。
每次处理一半，递归层数可达$log(N)$。

题解二（堆）

总所周知，小顶堆的根永远是值最小的节点，所以先将所有节点加入堆，然后在出堆k次即可得到结果。

这种方法实现简单，竞赛可以恰烂分。面试时如果面试官允许使用优先队列再用。

public int[] getLeastNumbers(int[] arr, int k) {
	if (k == 0 || arr == null || arr.length == 0) {
		return new int[]{};
	}

	PriorityQueue<Integer> pq = new PriorityQueue<>();
	for (int num : arr) {
		pq.offer(num);
	}
	int [] result = new int[k];
	for (int i = 0 ; i < result.length ; i++)
	{
		result[i] = pq.poll();
	}
	return result;
}

时间复杂度：$O(Nlog(N))$。
每次出堆，调整堆结构耗时$O(log(N)$，一共出堆$N$次。

空间复杂度：$O(N)$。
堆中需要存放$N$个元素。