手撕：TopK问题--1000个学生里面求成绩最高的10个

该题可归类为TopK问题，即求n个数据中的前k个数据，参考：

• 拜托，面试别再问我TopK了！！！

该问题需要了解一些前置知识，例如快排和堆排序，参考：

• 堆排序详细图解（通俗易懂）
• 排序——快速排序（Quick sort）

方案一：全局排序

复杂度为 $O(n\log n)$$

对n个数据进行快速排序，代码示例：

public static int[] topK_globalSort(int[] arr, int k) {
    // 全局排序，采用快排
    Arrays.sort(arr);
    // 输出topK
    return Arrays.copyOfRange(arr, arr.length - k, arr.length);
}

方案二：局部排序

复杂度为 $O(nk)$

采用冒泡排序，但只冒泡k次，剩下的n-k个数据不排序，代码示例：

public static int[] topK_localSort(int[] arr, int k) {
    // 局部排序，采用冒泡排序
    for (int i = 0; i < k; i++) {
        for (int j = i + 1; j < arr.length; j++) {
            if (arr[i] < arr[j]) {
                int temp = arr[i];
                arr[i] = arr[j];
                arr[j] = temp;
            }
        }
    }

    // 输出topK
    return Arrays.copyOfRange(arr, 0, k);
}

方案三：堆

复杂度为 $O(n\log k)$

构建大小为k的堆，遍历所有数据，最后堆内的k个数据就是目标数据，思路类似于局部排序，不过每次遍历时不要求堆内的数据有序，代码示例：

public static int[] topK_heap(int[] arr, int k) {
    // 通过PriorityQueue（底层实现为小顶堆）构建小顶堆
    PriorityQueue<Integer> heap = new PriorityQueue<>(k);
    for (int num : arr) {
        if (heap.size() < k) {
            // 堆中元素不足k个，直接添加
            heap.add(num);
        }

        if (heap.size() == k && num > heap.peek()) {
            // 堆中元素已满，且当前元素大于堆顶元素(堆中的最小值)，替换堆顶元素
            heap.poll();
            heap.add(num);
        }
    }

    // 输出堆中的元素
    return heap.stream().mapToInt(Integer::intValue).toArray();
}

方案四：减治法

复杂度为 $O(n)$

先区分一下分治法和减治法：

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

减治法的目的是尽可能少处理一些数据，应用到TopK问题中，就可以理解为快排的一种变体，只不过TopK是希望求出arr[1,n]中最大的k个数，那如果找到了第k大的数，做一次partition，就等于一次性找到最大的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大的元素即可；

代码示例：

public static int[] topK_decreaseAndConquer(int[] arr, int k) {
    // 减治法，目的是找到第k个元素
    int low = 0;
    int high = arr.length - 1;
    while (low < high) {
        int pivot = partition(arr, low, high);
        if (pivot == k - 1) {
            break;
        } else if (pivot < k - 1) {
            low = pivot + 1;
        } else {
            high = pivot - 1;
        }
    }

    // 输出topK
    return Arrays.copyOfRange(arr, 0, k);
}

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

测试：

public static void main(String[] args) {
    // 生成长度为n的数组, 值为随机数
    int n = 1000000;
    int[] arr = new int[n];
    for (int i = 0; i < arr.length; i++) {
        arr[i] = (int) (Math.random() * n);
    }
    int k = 100;

    // 打印各个方案的结果以及耗时
    long start = System.currentTimeMillis();
    //System.out.println("topK_globalSort: " + Arrays.toString(topK_globalSort(arr.clone(), k)));
    topK_globalSort(arr.clone(), k);
    System.out.println("topK_globalSort耗时: " + (System.currentTimeMillis() - start) + "ms");

    start = System.currentTimeMillis();
    //System.out.println("topK_localSort: " + Arrays.toString(topK_localSort(arr.clone(), k)));
    topK_localSort(arr.clone(), k);
    System.out.println("topK_localSort耗时: " + (System.currentTimeMillis() - start) + "ms");

    start = System.currentTimeMillis();
    //System.out.println("topK_heap: " + Arrays.toString(topK_heap(arr.clone(), k)));
    topK_heap(arr.clone(), k);
    System.out.println("topK_heap耗时: " + (System.currentTimeMillis() - start) + "ms");

    start = System.currentTimeMillis();
    //System.out.println("topK_decreaseAndConquer: " + Arrays.toString(topK_decreaseAndConquer(arr.clone(), k)));
    topK_decreaseAndConquer(arr.clone(), k);
    System.out.println("topK_decreaseAndConquer耗时: " + (System.currentTimeMillis() - start) + "ms");
}

结果为：

topK_globalSort耗时: 91ms
topK_localSort耗时: 74ms
topK_heap耗时: 13ms
topK_decreaseAndConquer耗时: 8ms

可以看到方案三和方案四的性能明显优于方案一和方案二。