Top K问题的解决方案

最新推荐文章于 2024-02-19 09:14:09 发布

努力的鑫

最新推荐文章于 2024-02-19 09:14:09 发布

阅读量699

点赞数

分类专栏：算法文章标签：算法

本文链接：https://blog.csdn.net/weixin_44163764/article/details/109597320

版权

算法专栏收录该内容

1 篇文章 0 订阅

订阅专栏

问题定义

Top K 问题指的是从一组无序数据中找到前 K 大（也可以是前 K 小）的数。这个问题是十分经典的问题，不论在面试中还是实际开发中，都非常典型。

解决方案

对于 TOP K 问题，常见的解决方案有以下三种：

排序
堆
快速选择

下面以剑指Offer 40.最小的K个数为例，实现上述的解决方案。

排序

排序是最容易想到的方法，将数据排序之后，取出最小的K个数即为所求。代码如下：

public int[] getLeastNumbers(int[] arr, int k) {
    Arrays.sort(arr);
    int[] res = new int[k];
    for (int i = 0; i < k; i++) {
        res[i] = arr[i];
    }
    return res;
}

Java语言中对于基本数据类型的数组排序使用的是快排，所以时间复杂度是O(nlogn)。
由于是对所有数据进行排序，所以时间复杂度太高，其实我们只需要 TOP K，所以能不能不全局排序，只对局部进行排序？其实，可以使用冒泡排序，这样只需要冒 K 次泡就可以得到 TOP K 了。代码如下：

public int[] getLeastNumbers(int[] arr, int k) {
    int[] res = new int[k];
    for (int i = 0; i < k; i++) {
        int index = i;
        for (int j = arr.length - 1; j > i; j--) {
            if (arr[j] < arr[index]) {
                index = j;
            }
        }
        int tmp = arr[i];
        arr[i] = arr[index];
        arr[index] = tmp;
        res[i] = arr[i];
    }
    return res;
}

时间复杂度是O(N * K)，将全局排序转化为局部排序，节省了资源，提高了效率。

堆

前面是排序的解决方案，那么能不能不排序就得到 TOP K呢？这时就可以借助堆来实现。
该问题是求前 K 小的数，那么可以创建一个最大堆，这个堆的容量为 K ，
当元素数量少于 K 个时，直接加入堆，等于 K 个时，将当前元素和堆顶元素进行比较。如果当前元素小于堆顶元素，堆顶元素出堆，当前元素入堆，这样比较完所有数据，堆中元素即为前 K 小的数。代码如下：

public int[] getLeastNumbers(int[] arr, int k) {
    if (arr.length == 0 || k == 0) {
        return new int[0];
    }
    PriorityQueue<Integer> heap = new PriorityQueue<>(((o1, o2) -> arr[o2] - arr[o1]));
    for (int i = 0; i < arr.length; i++) {
        if (heap.size() < k) {
            heap.offer(i);
        } else if (heap.comparator().compare(i, heap.peek()) > 0){
            heap.poll();
            heap.offer(i);
        }
    }
    int[] res = new int[k];
    int index = 0;
    for (Integer integer : heap) {
        res[index++] = arr[integer];
    }
    return res;
}

对于求前 K 大的情况，只需要创建小根堆，比较情况相反即可。

快速选择

快速排序是选择一个基准值，小于基准值的在左面，大于基准值的在右面。那么对于 TOP K 问题，只需要基准值的位置为 K 就可以了，这样左边的是前 K 小的，右边是前 K 大的。代码如下：

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

private int[] quickSearch(int[] array, int start, int end, int k) {
    int index = partiton(array, start, end);
    if (index == k - 1) {
        return Arrays.copyOf(array, k);
    }
    return index > k - 1 ? quickSearch(array, start, index - 1, k) : quickSearch(array, index + 1, end, k);
}

private int partiton(int[] array, int start, int end) {
    int pivot = start;
    swap(array, pivot, end);
    int index = start - 1;
    for (int i = start; i < end; i++) {
        if (array[i] < array[end]) {
            index++;
            if (index != i) {
                swap(array, index, i);
            }
        }
    }
    swap(array, index + 1, end);
    return index + 1;
}

private void swap(int[] array, int i, int j) {
    int tmp = array[i];
    array[i] = array[j];
    array[j] = tmp;
}