Top-k问题的解法

---

问题描述

所谓Top-K问题，指的是返回给定数据结构中最大/最小的K个元素，可能是本身的元素，也可能是运算结果中的元素。下面通过一简单的🌰来进行说明，如果解决这类问题。

---

最小的k个数

输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。

示例 1：

输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]

示例 2：

输入：arr = [0,1,2,1], k = 1
输出：[0]

限制：

• $0<=k<=arr.length<=10000$
• $0<=arr[i]<=10000$

---

1. 暴力排序法

解决上述问题最直接的方法进行对数组进行升序排序，然后返回结果前K个元素即可。例如：

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        Arrays.sort(arr);
        return Arrays.copyOfRange(arr, 0, k);
    }
}

复杂度分析：

• 时间复杂度：$O(N\log N)$，复杂度来自于排序过程，$N$为数据长度
• 空间复杂度：$O(\log N)$

---

2. 大根堆、小根堆

通过维护一个堆来保证元素之间的大小关系，例如大根堆中最大的元素位于堆顶，而小根堆中最小的元素位于堆顶。Java中相应的数据结构是优先级队列（PriorityQueue），通过传入不同的Comparator来设置堆的类型（默认为小根堆），initialCapacity用于设置堆的大小：

例如，使用PriorityQueue创建一个容量为K的大根堆：

Queue<Integer> bigHeap = new PriorityQueue(K,(k1, k2)-> k2 - k1 );

因此，该题的解法为：

class Solution {
   public int[] getLeastNumbers(int[] arr, int k) {
        if (arr.length == 0 || k == 0) {
            return new int[0];
        }

        Queue<Integer> heap = new PriorityQueue<>(k, (i1, i2) -> i2 - i1);
        for (int a : arr) {
            if (heap.isEmpty() || heap.size() < k) {
                heap.offer(a);
            } else if(heap.peek() > a){
                heap.poll();
                heap.offer(a);
            }
        }

        int[] result = new int[heap.size()];
        int index = 0;
        for (int e : heap) {
            result[index++] = e;
        }
        return result;
   }
}

---

3. 快速选择

快速排序的Java实现：

/** 
* @Author dyliang  
* @Date 2020/10/2 15:35  
* @Version 1.0  
*/
public class Sort {
    private static int[] arr = {3,5,1,6,2,9};

    public static void main(String[] args) {
        quickSort(arr, 0, arr.length - 1);
        System.out.println(Arrays.toString(arr));
    }

    private static void quickSort(int[] arr, int left, int right) {
        if (left < right) {
            int pivot = partition(arr, left, right);        // 将数组分为两部分
            quickSort(arr, left, pivot - 1);                   // 递归排序左子数组
            quickSort(arr, pivot + 1, right);                  // 递归排序右子数组
        }
    }

    private static int partition(int[] arr, int left, int right) {
        int pivot = arr[left];               // 枢轴记录
        while(left < right) {
            while (left < right && arr[right] >= pivot) --right;
            arr[left] = arr[right];           // 交换比枢轴小的记录到左端
            while (left < right && arr[left] <= pivot) ++left;
            arr[right] = arr[left];           // 交换比枢轴小的记录到右端
        }
        // 扫描完成，枢轴到位
        arr[left] = pivot;
        // 返回的是枢轴的位置
        return left;
    } 
}

快速选择类似于快速排序，它们都是分治思想的体现。对于快排来说，最为关键的一步是partition，在数组中选择一个元素作为pivot，将给定的数组按照元素和pivot大小的关系分为三个部分：小于pivot的元素部分、pivot和大于pivot的元素部分。递归进行此过程，递归结束时数组已有序。partition的过程需要遍历数组中全部的元素，因此，整个过程的时间复杂度为$O(N)$。

快速选择可以看做是快速排序的子过程，此时只关注某一部分元素的选择过程。例如，寻找数组中top-k个最小的元素，假设某一次partition后pivot元素所在的数组索引为$m$：

• $m==k$：那么，它之前的$k$的元素即所求结果，直接返回即可
• $m>k$：那么，最小的$k$个元素必定位于pivot的左侧，只需要对左侧数组递归进行partition
• $m<k$：那么，pivot最侧的$m$个元素是结果的一部分，还需对右侧数组递归进行partition，寻找右侧数组中最小的$k-m$个元素

对应的题目的解法为：

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if(k == 0 || arr.length == 0)   return new int[]{};
        if(arr.length <= k) return arr;
        return quickSearch(arr, 0, arr.length - 1, k);
    }

    public int[] quickSearch(int[] arr, int left, int right, int k) {
        int index = partition(arr, left, right);
        if(index == k) {
            return Arrays.copyOf(arr, index);
        }

        return index > k ? quickSearch(arr, left, index - 1, k) : quickSearch(arr, index + 1, right, k);
    }

    public int partition(int[] arr, int left, int right) {
        int pivot = arr[left]; 
        while(left < right) {
            while(left < right && arr[right] >= pivot) right--;
            arr[left] = arr[right];
            while(left < right && arr[left] <= pivot)  left++;
            arr[right] = arr[left];
        }
        arr[left] = pivot;
        return left;
    }
}

---

最接近原点的 K 个点

我们有一个由平面上的点组成的列表 points。需要从中找出 K 个距离原点 (0, 0) 最近的点。（这里，平面上两点之间的距离是欧几里德距离）你可以按任何顺序返回答案。除了点坐标的顺序之外，答案确保是唯一的。

示例 1：

输入：points = [[1,3],[-2,2]], K = 1
输出：[[-2,2]]
解释： 
(1, 3) 和原点之间的距离为 sqrt(10)，
(-2, 2) 和原点之间的距离为 sqrt(8)，
由于 sqrt(8) < sqrt(10)，(-2, 2) 离原点更近。
我们只需要距离原点最近的 K = 1 个点，所以答案就是 [[-2,2]]。

示例 2：

输入：points = [[3,3],[5,-1],[-2,4]], K = 2
输出：[[3,3],[-2,4]]
（答案 [[-2,4],[3,3]] 也会被接受。）

提示：

• $1<=K<=points.length<=10000$
• $-10000<points[i][0]<10000$
• $-10000<points[i][1]<10000$

这里直接给出快排的解法，套路是一样的。

class Solution {
    public int[][] kClosest(int[][] points, int K) {
        if(points.length == 0 || K == 0){
            return new int[][]{};
        }
        if(points.length <= K) 
            return points;

        return quickSearch(points, 0, points.length - 1, K);
    }

    public int[][] quickSearch(int[][] arr, int left, int right, int k){
        int index = partition(arr, left, right);
        if(index == k){
            return Arrays.copyOf(arr, index);
        }

        return index > k ? quickSearch(arr, left, index - 1, k) : quickSearch(arr, index + 1, right, k);
    }

    public int partition(int[][] arr, int left, int right){
        int[] pivot = arr[left];
        int distance = pivot[0] * pivot[0] + pivot[1] * pivot[1];
        while(left < right){
            while(left < right && arr[right][0]* arr[right][0] +  arr[right][1]* arr[right][1] >= distance) --right;
            arr[left] = arr[right];
            while(left < right && arr[left][0]* arr[left][0] +  arr[left][1]* arr[left][1] <= distance) ++left;
            arr[right] = arr[left];
        }

        arr[left] = pivot;
        return left;
    }
}