题目信息
示例1:
1 2
| 输入:arr = [3,2,1], k = 2 输出:[1,2] 或者 [2,1]
|
示例2:
1 2
| 输入:arr = [0,1,2,1], k = 1 输出:[0]
|
提示
1 2
| 1.0 <= k <= arr.length <= 10000 2.0 <= arr[i] <= 10000
|
解题思路
本题难点
多种解题方案,快排,大根堆,二叉搜索树,计数排序
具体思路
快速排序的思想。快排的划分函数每次执行完后都能将数组分成两个部分,小于等于分界值 pivot 的元素的都会被放到数组的左边,大于的都会被放到数组的右边,然后返回分界值的下标。
我们的目的是寻找最小的 k 个数。假设经过一次 partition 操作,分界值 pivot元素位于下标 j,也就是说,左侧的数组有 j 个元素,是原数组中最小的 j 个数。那么:
- k = j: 我们就找到了最小的 k 个数,就是左侧的数组;:
- k < j: 则最小的 k 个数一定都在左侧数组中,我们只需要对左侧数组递归地 parition 即可;
- k > j:则左侧数组中的 j 个数都属于最小的 k 个数,我们还需要在右侧数组中寻找最小的 k−j 个数,对右侧数组递归地 partition 即可。
代码
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41
| class Solution { public int[] getLeastNumbers(int[] arr, int k) { if(arr.length == 0 || k == 0){ return new int[0]; } return findKthSmallest(arr,0,arr.length-1,k-1); }
public int[] findKthSmallest(int[] arr,int l ,int h , int k){ int j = partition(arr,l,h); if(j == k){ return Arrays.copyOf(arr, j + 1); } return j > k ? findKthSmallest(arr,l,j-1,k):findKthSmallest(arr,j+1,h,k); }
public int partition(int[] arr,int l,int h){ int privot = arr[l]; int i = l,j = h + 1; while(true){ while(i != h && arr[++i] < privot); while(j != l && arr[--j] > privot); if(i >= j){ break; } swap(arr,i,j); } swap(arr,l,j); return j; } public void swap(int[] arr,int i,int j){ int temp = arr[i]; arr[i] = arr[j]; arr[j] = temp; } }
|
复杂度分析:
- 时间复杂度 O(N) : 找下标为k的元素,第一次切分的时候需要遍历整个数组 (0 ~ n) 找到了下标是 j 的元素,假如 k 比 j 小的话,那么我们下次切分只要遍历数组 (0~k-1)的元素就行啦,总之可以看作每次调用 partition 遍历的元素数目都是上一次遍历的 1/2,因此时间复杂度是 N + N/2 + N/4 + … + N/N = 2N, 因此时间复杂度是 O(N)。
- 空间复杂度 O(logN) : 递归调用的期望深度为O(logN),每层需要的空间为 O(1),只有常数个变量。
其他优秀解答
解题思路
本题是求前 K 小,因此用一个容量为 K 的大根堆,每次 poll 出最大的数,那堆中保留的就是前 K 小啦(注意不是小根堆!小根堆的话需要把全部的元素都入堆,那是 O(NlogN),就不是 O(NlogK))这个方法比快排慢。
代码
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29
|
class Solution { public int[] getLeastNumbers(int[] arr, int k) { if (k == 0 || arr.length == 0) { return new int[0]; } Queue<Integer> pq = new PriorityQueue<>((v1, v2) -> v2 - v1); for (int num: arr) { if (pq.size() < k) { pq.offer(num); } else if (num < pq.peek()) { pq.poll(); pq.offer(num); } } int[] res = new int[pq.size()]; int idx = 0; for(int num: pq) { res[idx++] = num; } return res; } }
|
解题思路
BST 相对于前两种方法没那么常见,但是也很简单,和大根堆的思路差不多,与前两种方法相比,BST 有一个好处是求得的前K大的数字是有序的。
代码
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43
| class Solution { public int[] getLeastNumbers(int[] arr, int k) { if (k == 0 || arr.length == 0) { return new int[0]; } TreeMap<Integer, Integer> map = new TreeMap<>(); int cnt = 0; for (int num: arr) { if (cnt < k) { map.put(num, map.getOrDefault(num, 0) + 1); cnt++; continue; } Map.Entry<Integer, Integer> entry = map.lastEntry(); if (entry.getKey() > num) { map.put(num, map.getOrDefault(num, 0) + 1); if (entry.getValue() == 1) { map.pollLastEntry(); } else { map.put(entry.getKey(), entry.getValue() - 1); } } }
int[] res = new int[k]; int idx = 0; for (Map.Entry<Integer, Integer> entry: map.entrySet()) { int freq = entry.getValue(); while (freq-- > 0) { res[idx++] = entry.getKey(); } } return res; } }
|
解题思路
数据范围有限时直接计数排序就行了:O(N)
代码
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24
| class Solution { public int[] getLeastNumbers(int[] arr, int k) { if (k == 0 || arr.length == 0) { return new int[0]; } int[] counter = new int[10001]; for (int num: arr) { counter[num]++; } int[] res = new int[k]; int idx = 0; for (int num = 0; num < counter.length; num++) { while (counter[num]-- > 0 && idx < k) { res[idx++] = num; } if (idx == k) { break; } } return res; } }
|