TopK问题的两种常见解法

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;
      }
  }