【算法题刷完这遍该会了】最小 K 个数 - TopK 问题

本文链接：https://blog.youkuaiyun.com/qq_44740757/article/details/147038776

经典 TopK 问题，面试高频问题，前两天某团二面刚被考过。
力扣链接 - 最小 K 个数

文章目录

题目描述

设计一个算法，找出数组中最小的 k 个数。以任意顺序返回这 k 个数即可。

示例 1:

输入: arr = [1,3,5,7,2,4,6,8], k = 4
输出: [1,2,3,4]

示例 2:

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

示例 3:

输入: arr = [1,2,4,5], k = 3
输出: [1,2,4]

提示：

0 <= len (arr) <= 100000
0 <= k <= min (100000, len (arr))

解题思路

这道题有多种解法，我们从简单到复杂逐步分析：
拿到题目首先应该都能想到要排序，直接 Arrays.sort 是 O(nlogn) 没毛病。

解法一：排序（最直观）

最简单的方法是先对数组排序，然后取前 k 个元素。

class Solution {
    public int[] smallestK(int[] arr, int k) {
        if (k == 0) return new int[0];
        Arrays.sort(arr);
        return Arrays.copyOfRange(arr, 0, k);
    }
}

复杂度分析：

时间复杂度：O (nlogn)，排序的时间复杂度
空间复杂度：O (logn)，排序需要的栈空间

不得不说 Arrays.sort() 优化得真好

解法二：优先队列之小顶堆（直观）

基本思路依然是全排序然后取前 K 个，但是这里使用堆排序，有个好处是并不需要完全排序，建堆后依次取出 K 个堆顶元素就可以提前结束。

class Solution {
    public int[] smallestK(int[] arr, int k) {
        if (k == 0) return new int[0];
        // 小顶堆
        PriorityQueue<Integer> pq = new PriorityQueue<>();
        
        // 将所有元素入堆
        for (int num : arr) {
            pq.offer(num);
        }
        
        // 取出k个最小元素
        int[] res = new int[k];
        for (int i = 0; i < k; i++) {
            res[i] = pq.poll();
        }
        return res;
    }
}

复杂度分析：

时间复杂度：O(nlogn)，建堆的时间复杂度
空间复杂度：O(n)，优先队列的大小

解法三：优先队列之大顶堆（优秀）

更加充分利用堆排序的特点。直接维护一个代表结果集的堆，充分利用大顶堆堆顶永远是当前堆中最大的元素这个特点。先将前 k 个元素压入堆中，此时堆顶的最大元素就是最需要淘汰和替换的，从 K+1 元素继续遍历，只要有比堆顶小的，替换之。
形象化记忆：堆里都是牛马，堆顶是年龄最大的，只要招到更年轻的牛马就替换之（蚌埠住了）。
使用优先队列（大顶堆）维护 k 个最小的元素。当遇到更小的元素时，替换堆顶元素。

class Solution {
    public int[] smallestK(int[] arr, int k) {
        if (k == 0) return new int[0];
        // 大顶堆
        PriorityQueue<Integer> pq = new PriorityQueue<>((a, b) -> b - a);
        
        for (int num : arr) {
            if (pq.size() < k) { // 先随便放 k 个元素进去
                pq.offer(num);
            } else if (num < pq.peek()) { // 有更年轻的牛马
            	// 替换堆顶元素
                pq.poll();
                pq.offer(num);
            }
        }
        
        int[] res = new int[k];
        for (int i = 0; i < k; i++) {
            res[i] = pq.poll();
        }
        return res;
    }
}

复杂度分析：

时间复杂度：O (nlogk)， K 个元素的建堆时间
空间复杂度：O (k)，优先队列的大小

解法四：快速选择（最优解）

使用快速排序的思想，通过 partition 操作找到第 k 小的元素，其左边的元素就是最小的 k 个数。
快速排序如果不熟悉可以参考 手撕快速排序

核心思想类比：

还是那个例子，要找试卷中第 k 低分的那张试卷。
你不需要把所有试卷都完全排序，也不需要先找到最低分的，再找第二低分的…快速选择的做法是：

选个“基准”(Pivot)：随机选一张试卷作为基准值（如果你运气不好，选到的是总是最高分，那就要尴尬了，所以选择基准这里是有优化空间的）。
分区(Partition)：快速把其他试卷分成两堆：分数 ≤ 基准的，和分数 > 基准的。基准值本身也会被放到它最终应该在的位置，假设这个位置的索引是 pivotIndex。
判断与缩小范围：

如果 pivotIndex 正好是 k-1（因为索引从0开始，第k小的元素在索引k-1处），那么基准左边的所有元素（包括基准自己）就是最小的 k 个数，找到了！
如果 pivotIndex < k-1，说明第 k 小的数肯定在基准右边那堆更大的数里。我们只需要在右半部分继续寻找（递归）。
如果 pivotIndex > k-1，说明第 k 小的数在基准左边那堆更小的数里。我们只需要在左半部分继续寻找（递归）。

关键优势：

每次分区后，我们都能排除掉一部分数据（要么左半部分，要么右半部分），不用再管它们了。平均情况下，每次都能排除掉大约一半的数据，所以总的操作次数大约是 n + n/2 + n/4 + … ≈ 2n，因此平均时间复杂度是 O(n)。

复杂度

平均时间： O(n) - 非常高效。
最坏时间： O(n²) - 如果每次选的基准都恰好是当前区间的最大或最小值（例如数组已排序），分区就会很不均匀。可以通过随机选基准或“三数取中法”来大大降低发生最坏情况的概率。
空间： O(log n) - 主要是递归调用栈的深度。

注意这里不能保证有序输出前 K 小的数。题目如果变式为有序输出最小 K 个数则不考虑此方法。

class Solution {
	public int[] smallestK(int[] arr, int k) {
		if (k == 0) return new int[0];
		quickSelect(arr, 0, arr.length - 1, k);
		return Arrays.copyOfRange(arr, 0, k);
	}
	
	private void quickSelect(int[] arr, int left, int right, int k) {
		if (left >= right) return;
		
		int pivot = partition(arr, left, right);
		if (pivot == k) {
			return;
		}
		
		if (pivot < k) {
			quickSelect(arr, pivot + 1, right, k);
		} else {
			quickSelect(arr, left, pivot - 1, k);
		}
	}
	
	private int partition(int[] arr, int left, int right) {
		// 还是三数取中法优化最差情况下的效率
		int mid = left + (right - left) / 2;
		// 三步比较确保 left <= mid <= right
		if (arr[left] > arr[mid]) swap(arr, left, mid); // 换完 mid > left
		if (arr[mid] > arr[right]) swap(arr, mid, right); // 换完 right > mid, right > left
		if (arr[left] > arr[mid]) swap(arr, left, mid); // 换完 right > mid > left
		
		// 取出 pivot 作为基准值，同时将基准换到最左侧
		int pivot = arr[mid];
		swap(arr, mid, left);
		
		// 相向双指针，交换后左侧均小于等于 pivot，右侧大于等于 pivot
		int i = left, j = right;
		while(i < j) {
			while(i < j && arr[j] >= pivot) j--; // 注意 >=
			while(i < j && arr[i] <= pivot) i++; // 注意 <=
			if (arr[i] > arr[j]) swap(arr, i, j);
		}
		
		swap(arr, i, left);
		return i;
	}
	
	private void swap(int[] arr, int i, int j) {
		int temp = arr[i];
		arr[i] = arr[j];
		arr[j] = temp;
	}
}