寻找第k小元素 前k小元素 select_k

问题

对于海量数据到处理经常会涉及到 topK 问题。在设计数据结构和算法的时候，主要需要考虑的应该是当前算法（包括数据结构）跟给定情境（比如数据量级、数据类型）的适配程度，和当前问题最核心的瓶颈（如降低时间复杂度，还是降低空间复杂度）是什么。

首先，我们来举几个常见的 topK 问题的例子：

给定 100 个 int 数字，在其中找出最大的 10 个；
给定 10 亿个 int 数字，在其中找出最大的 10 个（这 10 个数字可以无序）；
给定 10 亿个 int 数字，在其中找出最大的 10 个（这 10 个数字依次排序）；
给定 10 亿个不重复的 int 数字，在其中找出最大的 10 个；
给定 10 个数组，每个数组中有 1 亿个 int 数字，在其中找出最大的 10 个；
给定 10 亿个 string 类型的数字，在其中找出最大的 10 个（仅需要查 1 次）；
给定 10 亿个 string 类型的数字，在其中找出最大的 k 个（需要反复多次查询，其中 k 是一个随机数字）。

上面这些问题看起来很相似，但是解决的方式却千差万别。稍有不慎，就可能使得 topK 问题成为系统的瓶颈。不过也不用太担心，接下来我会总结几种常见的解决思路，遇到问题的时候，大家把这些基础思路融会贯通并且杂糅组合，即可做到见招拆招。

最基本思路

将N个数进行完全排序，从中选出排在前K的元素即为所求。有了这个思路，我们可以选择相应的排序算法进行处理，目前来看快速排序，堆排序和归并排序都能达到O(NlogN)的时间复杂度。

优先队列

选择其中前K个数作为数据池，后面的N-K个数与这K个数进行比较，若小于/大于其中的任何一个数，则进行替换。这种思路的算法复杂度是O(N*K).
剩余的N-K个数与前面K个数比较的时候，是顺序比较的，算法复杂度是O(N*K)。

小根堆

堆排序是通过维护大顶堆或者小顶堆来实现的。堆排序法来解决N个数中的TopK的思路是：

1. 先随机取出N个数中的K个数，将这K个数构造为小顶堆，那么堆顶的数肯定就是这K个数中最小的数了
2. 然后再将剩下的N-K个数与堆顶进行比较，如果大于堆顶，那么说明该数有机会成为TopK，就更新堆顶为该数，此时由于小顶堆的性质可能被破坏，就还需要调整堆；

//添加每个数到小根堆，小根堆堆化
//如果 num < 堆顶(最小元素)，会堆化到堆顶，如果 num > 堆顶，会堆化到堆顶之下
queue.add(num);		
if(queue.size()>k) queue.poll();    //把 堆中最小的 poll 了

4. 否则说明这个数最多只能成为Top K+1，因此就不用管它 。
5. 然后就将下一个数与当前堆顶的数作比较，根据大小关系如上面所述方法进行操作，直到N-K个数都遍历完，此时还在堆中的K个数就是TopK了。

对K个元素进行建堆，时间复杂度为O(K)；
然后对剩下的N-K个数对堆顶进行比较及更新，最好情况下当然是都不需要调整了，那么时间复杂度就只是遍历这N-K个数的O(N-K)，这样总体的时间复杂度就是O(N)
而在最坏情况下，N-K个数都需要更新堆顶，每次调整堆的时间复杂度为logK，因此此时时间复杂度就是NlogK了，总的时间复杂度就是O(K)+O(NlogK)≈O(NlogK)

一般来说企业中都采用该策略处理topK问题，因为该算法不需要一次将原数组中的内容全部加载到内存中，而这正是海量数据处理必然会面临的一个关卡。

快速选择算法

快速选择算法 - Quick select

利用快速排序的partition分划函数找到分划位置K，则其前面的内容即为所求。该算法是一种非常有效的处理方式，时间复杂度是O(N)（证明可以参考算法导论书籍）。对于能一次加载到内存中的数组，该策略非常优秀。

“冒泡排序”获取TopK

可以避免对所有数据进行排序，只排序部分；
冒泡排序是每一轮排序都会获得一个最大值，则K轮排序即可获得TopK。
时间复杂度O(N*K)

分治法 + 快速选择算法

1. 比如有10亿的数据，找处Top1000，我们先将10亿的数据分成1000份，每份100万条数据。
2. 在每一份中找出对应的Top 1000，整合到一个数组中，得到100万条数据，这样过滤掉了99%的数据。
3. 使用快速排序对这100万条数据进行”一轮“排序，一轮排序之后指针的位置指向的数字假设为S，会将数组分为两部分，一部分大于S记作Si，一部分小于S记作Sj。
4. 如果Si元素个数大于1000，我们对Si数组再进行一轮排序，再次将Si分成了Sii和Sjj。如果Si的元素小于1000，则我们需要在Sj中获取1000-count(Si)个元素的，也就是对Sj进行排序
5. 如此递归下去即可获得Top-K。

总结

使用最大最小堆。求最大的数用最小堆，求最小的数用最大堆。     O(K)+O(NlogK)≈O(NlogK)
		一次堆化的时间复杂度为 O(logK)

Quick Select算法。使用类似快排的思路，根据pivot划分数组。  O(N)	
使用排序方法，排序后再寻找top K元素。  				     O(NlogN)
使用选择排序的思想，对前K个元素部分排序。

优先队列。前K个数作为数据池，后面的N-K个数与这K个数进行比较，若小于/大于其中的任何一个数，则进行替换			O(N*K)

将1000.....个数分成m组，每组寻找top K个数，得到m×K个数，在这m×k个数分治。

leetcode之TopK（快排/堆排序）

• 快排思想
  并不需要把全部都排序好，只用分组快排，快排其实是把小于基准数的所有数放在左边，大于的数都放在右边，只要找到这个基准数在快排后的下标，若下标<k-1，则将左边那组继续快排，若>k-1，则将右边那组快排。

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }
        // 最后一个参数表示我们要找的是下标为k-1的数
        return quickSearch(arr, 0, arr.length - 1, k - 1);
    }

    private int[] quickSearch(int[] nums, int lo, int hi, int k) {
        int j = partition(nums, lo, hi);
        if (j == k) {
            return Arrays.copyOf(nums, j + 1);
        }
        // 否则根据下标j与k的大小关系来决定继续切分左段还是右段。
        return j > k? quickSearch(nums, lo, j - 1, k): quickSearch(nums, j + 1, hi, k);
    }

    // 快排切分
    private int partition(int[] nums, int lo, int hi) {
        int v = nums[lo];
        int i = lo, j = hi + 1;
        while (true) {
            while (++i <= hi && nums[i] < v);
            while (--j >= lo && nums[j] > v);
            if (i >= j) {
                break;
            }
            int t = nums[j];
            nums[j] = nums[i];
            nums[i] = t;
        }
        nums[lo] = nums[j];
        nums[j] = v;
        return j;
    }
}

• 堆排序
  使用Java自带的PriorityQueue的方法，只需要k个值，每次往堆中添加时，若不满足条件则不添加。

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
      if(arr.length==0||k==0) return new int[0];
        //堆默认是小跟堆，(i1,i2)->Integer.compare(i2,i1)变成大根堆
        Queue<Integer> heap=new PriorityQueue<>(k,(i1,i2)->Integer.compare(i2,i1));
        for (int e:
             arr) {
            if (heap.isEmpty() || heap.size()<k || e<heap.peek()){
                heap.offer(e);
            }
            if(heap.size()>k){
                heap.poll();
            }
        }
        int[] nums=new int[heap.size()];
        int j=0;
        for (int e:
            heap ) {
            nums[j++]=e;
        }
        return nums;
    }
}

个人博客请访问 http://www.x0100.top 

TopK，是问得比较多的几个问题之一，到底有几种方法，这些方案里蕴含的优化思路究竟是怎么样的，今天和大家聊一聊。

问题描述：

从arr[1, n]这n个数中，找出最大的k个数，这就是经典的TopK问题。

栗子：

从arr[1, 12]={5,3,7,1,8,2,9,4,7,2,6,6} 这n=12个数中，找出最大的k=5个。

 

一、排序

排序是最容易想到的方法，将n个数排序之后，取出最大的k个，即为所得。

 

伪代码：

sort(arr, 1, n);

return arr[1, k];

 

时间复杂度：O(n*lg(n))
 

分析：明明只需要TopK，却将全局都排序了，这也是这个方法复杂度非常高的原因。那能不能不全局排序，而只局部排序呢？这就引出了第二个优化方法。

 

二、局部排序

不再全局排序，只对最大的k个排序。

冒泡是一个很常见的排序方法，每冒一个泡，找出最大值，冒k个泡，就得到TopK。

 

伪代码：

for(i=1 to k){

         bubble_find_max(arr,i);

}

return arr[1, k];

 

时间复杂度：O(n*k)

 

分析：冒泡，将全局排序优化为了局部排序，非TopK的元素是不需要排序的，节省了计算资源。不少朋友会想到，需求是TopK，是不是这最大的k个元素也不需要排序呢？这就引出了第三个优化方法。

 

三、堆

思路：只找到TopK，不排序TopK。

先用前k个元素生成一个小顶堆，这个小顶堆用于存储，当前最大的k个元素。

 

接着，从第k+1个元素开始扫描，和堆顶（堆中最小的元素）比较，如果被扫描的元素大于堆顶，则替换堆顶的元素，并调整堆，以保证堆内的k个元素，总是当前最大的k个元素。

 

直到，扫描完所有n-k个元素，最终堆中的k个元素，就是猥琐求的TopK。

 

伪代码：

heap[k] = make_heap(arr[1, k]);

for(i=k+1 to n){

         adjust_heap(heep[k],arr[i]);

}

return heap[k];

 

时间复杂度：O(n*lg(k))

画外音：n个元素扫一遍，假设运气很差，每次都入堆调整，调整时间复杂度为堆的高度，即lg(k)，故整体时间复杂度是n*lg(k)。

 

分析：堆，将冒泡的TopK排序优化为了TopK不排序，节省了计算资源。堆，是求TopK的经典算法，那还有没有更快的方案呢？

 

四、随机选择

随机选择算在是《算法导论》中一个经典的算法，其时间复杂度为O(n)，是一个线性复杂度的方法。

 

这个方法并不是所有同学都知道，为了将算法讲透，先聊一些前序知识，一个所有程序员都应该烂熟于胸的经典算法：快速排序。

画外音：

（1）如果有朋友说，“不知道快速排序，也不妨碍我写业务代码呀”…额...

（2）除非校招，我在面试过程中从不问快速排序，默认所有工程师都知道；

 

其伪代码是：

void quick_sort(int[]arr, int low, inthigh){

         if(low== high) return;

         int i = partition(arr, low, high);

         quick_sort(arr, low, i-1);

         quick_sort(arr, i+1, high);

}

 

其核心算法思想是，分治法。

 

分治法（Divide&Conquer），把一个大的问题，转化为若干个子问题（Divide），每个子问题“都”解决，大的问题便随之解决（Conquer）。这里的关键词是“都”。从伪代码里可以看到，快速排序递归时，先通过partition把数组分隔为两个部分，两个部分“都”要再次递归。

 

分治法有一个特例，叫减治法。

 

减治法（Reduce&Conquer），把一个大的问题，转化为若干个子问题（Reduce），这些子问题中“只”解决一个，大的问题便随之解决（Conquer）。这里的关键词是“只”。

 

二分查找binary_search，BS，是一个典型的运用减治法思想的算法，其伪代码是：

int BS(int[]arr, int low, inthigh, int target){

         if(low> high) return -1;

         mid= (low+high)/2;

         if(arr[mid]== target) return mid;

         if(arr[mid]> target)

                   return BS(arr, low, mid-1, target);

         else

                   return BS(arr, mid+1, high, target);

}

 

从伪代码可以看到，二分查找，一个大的问题，可以用一个mid元素，分成左半区，右半区两个子问题。而左右两个子问题，只需要解决其中一个，递归一次，就能够解决二分查找全局的问题。

 

通过分治法与减治法的描述，可以发现，分治法的复杂度一般来说是大于减治法的：

快速排序：O(n*lg(n))

二分查找：O(lg(n))

 

话题收回来，快速排序的核心是：

i = partition(arr, low, high);

 

这个partition是干嘛的呢？

顾名思义，partition会把整体分为两个部分。

更具体的，会用数组arr中的一个元素（默认是第一个元素t=arr[low]）为划分依据，将数据arr[low, high]划分成左右两个子数组：

• 左半部分，都比t大

• 右半部分，都比t小

• 中间位置i是划分元素

以上述TopK的数组为例，先用第一个元素t=arr[low]为划分依据，扫描一遍数组，把数组分成了两个半区：

• 左半区比t大

• 右半区比t小

• 中间是t

partition返回的是t最终的位置i。

 

很容易知道，partition的时间复杂度是O(n)。

画外音：把整个数组扫一遍，比t大的放左边，比t小的放右边，最后t放在中间N[i]。

 

partition和TopK问题有什么关系呢？

TopK是希望求出arr[1,n]中最大的k个数，那如果找到了第k大的数，做一次partition，不就一次性找到最大的k个数了么？

画外音：即partition后左半区的k个数。

 

问题变成了arr[1, n]中找到第k大的数。

 

再回过头来看看第一次partition，划分之后：

i = partition(arr, 1, n);

• 如果i大于k，则说明arr[i]左边的元素都大于k，于是只递归arr[1, i-1]里第k大的元素即可；

• 如果i小于k，则说明说明第k大的元素在arr[i]的右边，于是只递归arr[i+1, n]里第k-i大的元素即可；

画外音：这一段非常重要，多读几遍。

 

这就是随机选择算法randomized_select，RS，其伪代码如下：

int RS(arr, low, high, k){

  if(low== high) return arr[low];

  i= partition(arr, low, high);

  temp= i-low; //数组前半部分元素个数

  if(temp>=k)

      return RS(arr, low, i-1, k); //求前半部分第k大

  else

      return RS(arr, i+1, high, k-i); //求后半部分第k-i大

}

 

这是一个典型的减治算法，递归内的两个分支，最终只会执行一个，它的时间复杂度是O(n)。

 

再次强调一下：

• 分治法，大问题分解为小问题，小问题都要递归各个分支，例如：快速排序

• 减治法，大问题分解为小问题，小问题只要递归一个分支，例如：二分查找，随机选择

 

通过随机选择（randomized_select），找到arr[1, n]中第k大的数，再进行一次partition，就能得到TopK的结果。

 

五、总结

TopK，不难；其思路优化过程，不简单：

• 全局排序，O(n*lg(n))

• 局部排序，只排序TopK个数，O(n*k)

• 堆，TopK个数也不排序了，O(n*lg(k))

• 分治法，每个分支“都要”递归，例如：快速排序，O(n*lg(n))

• 减治法，“只要”递归一个分支，例如：二分查找O(lg(n))，随机选择O(n)

• TopK的另一个解法：随机选择+partition

 

知其然，知其所以然。

思路比结论重要。

更多精彩内容扫描下方二维码进入网站。。。。。

关注微信公众号。。。。。