TopK数的解法汇总

寻找第K大

https://www.nowcoder.com/practice/e016ad9b7f0b45048c58a9f27ba618bf?tpId=190&&tqId=35209&rp=1&ru=/ta/job-code-high-rd&qru=/ta/job-code-high-rd/question-ranking

TopK问题，不管是求前K大/前K小/第K大/第K小等，都有下面四种解法：

O(N)：用快排变形最最最高效解决TopK问题
O(NlogK)：大根堆（前K小）/小根堆（前K大）
O(NlogK)：二叉搜索树
O(N): 对于数据范围有限的情况可以直接计数排序O(N)高效解决

我们上一道例题：
图片说明

方法一快排：
时间复杂度为O(n)

class Solution {
  public int[] getLeastNumbers(int[] arr, int k) {
      if(k == 0 || arr.length == 0){
          return new int[0];
      }
      //因为我们要的是前k个数，对应下标为k-1
      return quickSort(arr, 0, arr.length - 1, k - 1);
  }

  public static int[] quickSort(int[] arr, int L, int R, int k){
      int i = partition(arr, L, R);
      //恰好相等，说明左边已经有k个较小的数了
      if(i == k){
          return Arrays.copyOf(arr, i+1);
      }

      return i > k ? quickSort(arr, L, i-1, k) : quickSort(arr, i+1, R, k);
  }

  //填坑的过程
  public static int partition(int[] arr, int L, int R){
      int base = arr[L];
      int l = L, r = R;
      while(l < r){
          while(l < r && arr[r] >= base){
              r--;
          }
          if(l < r){
              arr[l] = arr[r];
              l++;
          }
          while(l < r && arr[l] < base){
              l++;
          }
          if(l < r){
              arr[r] = arr[l];
              r--;
          }

      }
      //呜呜呜一定要记得填回去啊 
      arr[l] = base;
      return l;
  }
}

因为我们是要找下标为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(NlongN)
用堆时间复杂度会比快排要慢很多，但是Java提供了现成的PriorityQueue（默认小根堆），索引代码实现起来很简单。
本题是求前K小，因此用一个容量为K的大根堆（每次poll出最大的数，那堆中保留的就是前K个小的数）

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){
          //构造一个k大的大根堆  O(N)
          if(pq.size() < k){
              pq.offer(num); 
          }else{ 
              if(num < pq.peek()){
                  pq.poll(); //超过K,调整堆 O(NlongN)
                  pq.offer(num);
              }
          }
      }
      int[] res = new int[pq.size()];
      int index = 0;
      for(int num : pq){
          res[index++] = num; 
      }

      return res;
  }
}

方法三 BST：
时间复杂度 O(NlogN)
因为有重复的数字，所以用的是TreeMap而不是TreeSet（有的语言的标准库自带TreeMultiset，也是可以的）。TreeMap的key是数字，value是该数字的个数。我们遍历数组中的数字，维护一个数字总个数为K的TreeMap，每遍历一个元素：

TopK数的解法汇总

全站热榜