Top K 问题一网打尽

Top K 问题

Top K 问题，指的是在一个无序数组中，找前 K 大的元素或者第 K 大的元素，这两个问题其实是一个问题，有着共同的解法.
本文将针对【找出第K大元素】这个问题给出 4 中不同的解法。

解决方案

1 遍历 + 插入排序，O(nk)

我们可以想到最暴力简单的方式就是把 K 个最大的元素保存在一个有序数组中，每次从原数组遍历拿一个元素，用插入排序的方式去更新这个有序数组，最后，这个数组就是 Top K 的元素。

import java.lang.reflect.Array;
import java.util.ArrayList;

/**
 * Top k 问题代码
 */
public class Main {
    public static void main(String[] args) {
        System.out.println(topK(new int[]{7,9,4,5,8,8,1,0,1},7));
    }

    /**
     * 遍历+插入排序
     * @param num
     */
    public static int topK(int[] num, int k){
        int n = num.length;
        if (k > n) {
            return 0;
        }
        if (k <= 0) return 0;
        int[] a = new int[k + 1];//插入排序的数组
        a[0] = num[0];           // 先插入一个
        int curLength = 1;       // 插入排序数组的元素数目

        for (int i = 1; i < n; i++) { //遍历数组
            System.out.println();
            int insert;//下一个要插入的位置
            for (insert = curLength - 1; insert >= 0; insert--) {
                if (num[i] <= a[insert]) {
                    break;
                }else {
                    a[insert + 1] = a[insert];//插入
                }
            }
            a[insert+1] = num[i];
            if (curLength < k) {    //最多保留K个
                curLength++;
            }

        }
        return a[k - 1];

    }
}

2 排序，O(nlogn)

也可用堆数组进行排序，然后取 Top K 个：

import java.lang.reflect.Array;
import java.util.ArrayList;
import java.util.Arrays;

/**
 * Top k 问题代码
 */
public class Main {
    public static void main(String[] args) {
        System.out.println(topK(new int[]{7,9,4,5,8,8,1,0,1},9));
    }

    public static int topK(int[] num, int k){
        int n = num.length;
        if (k > n) {
            return 0;
        }
        if (k <= 0) return 0;
        Arrays.sort(num);
        return num[n - k];

    }
}

3 最小堆，O(nlogK)

为了保存Top K个元素，也可以用一个最小堆来保存，原理和插入排序一样，但是时间更快。

至于为什么是最小堆不是最大堆，因为最小堆的堆顶是堆中最小的拿个元素，这样的话有更大的元素进来，我们只需要淘汰堆顶就可以了。

import java.lang.reflect.Array;
import java.util.ArrayList;
import java.util.PriorityQueue;

/**
 * Top k 问题代码
 */
public class Main {
    public static void main(String[] args) {
        System.out.println(topK(new int[]{7,9,4,5,8,8,1,0,1},6));
    }

    /**
     * 遍历+插入排序
     * @param num
     */
    public static int topK(int[] num, int k){
        int n = num.length;
        if (k > n) {
            return 0;
        }
        if (k <= 0) return 0;

        PriorityQueue<Integer> minQ = new PriorityQueue<>((o1, o2) -> o1.compareTo(o2));
        for (int i = 0; i < n; i++) { //遍历数组
            minQ.add(num[i]);
            if (minQ.size() > k) {//淘汰堆顶
                minQ.remove();
            }
        }
        return minQ.peek();

    }
}

4 快速选择， O(n)

注意我们的排序方法，排序对所有元素进行排序，但是我们并不需要这个信息，我们只需要 K 个最大的元素，不需要排序。排序是多余的，我们只需要筛选 K 个最大的就行了。

类似于快排的思路，如图所示，每一趟，我们把数组进行划分成两个部分，前一部分比后一部分所有的元素都大，这样，我们只需要判断第K个元素在哪一部分就行了，然后递归地继续划分另一部分。
这样，平均复杂度是 O(n)。

import java.lang.reflect.Array;
import java.util.ArrayList;
import java.util.PriorityQueue;

/**
 * Top k 问题代码
 */
public class Main {
    public static void main(String[] args) {
        System.out.println(topK(new int[]{1,11,7,9,4,5,8,8,1,0,1},1));
    }

    /**
     * 遍历+插入排序
     * @param num
     */
    public static int topK(int[] num, int k){
        int n = num.length;
        if (k > n) {
            return 0;
        }
        if (k <= 0) return 0;

       partition(num, 0, n - 1, k);
       return num[k-1];

    }
    private static void partition(int[] num, int l, int r, int k) {
        if (l >= r) {
            return;
        }
        int temp = num[l];
        int left = l;
        int right = r;
        while (left < right) {//类似于快速排序的划分,把大于temp的放在左边
            while (left < right) {
                if (num[right] >= temp) {
                    num[left++] = num[right];
                    break;
                }
                right--;
            }
            while (left < right) {
                if (num[left] < temp) {
                    num[right--] = num[left];
                    break;
                }
                left++;
            }
        }
        num[left] = temp;

        if (left - l + 1 == k){//正好是第k个
            return;
        }else if (left - l + 1 < k) {//在右边
            partition(num,left + 1, r, k - left + l - 1);
        }else {//在左边
            partition(num,l, left - 1, k);
        }
    }
}