5. TopK问题
-
解法1:改进的快排 O(n)
public int[] getLeastNumbers(int[] arr, int k) { if(arr.length<=k){return arr;} sort(arr,0,arr.length-1,k); return Arrays.copyOfRange(arr,0,k); } public void sort(int[] arr,int start, int end, int k){ if(start>=end){return;} int flag = arr[start]; int l=start,r=end; while(l<r){ while(l<end&&arr[l]<=flag){ l++; } while(r>start&&arr[r]>=flag){ r--; } if(l>=r){ break; } int temp = arr[l]; arr[l] = arr[r]; arr[r] =temp; } arr[start]=arr[r]; arr[r]=flag; if(r>=k-1){ sort(arr,start,r,k); }else{ sort(arr,l,end,k); } }
-
解法二:大顶堆(O(nlogK))
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; } }