top-k问题详解——通过堆解决高频面试题

最新推荐文章于 2022-10-05 13:22:12 发布

陈亦康

最新推荐文章于 2022-10-05 13:22:12 发布

阅读量295

点赞数 2

文章标签：面试题 top-k问题堆优先级队列

本文链接：https://blog.csdn.net/CYK_byte/article/details/126207184

版权

优先队列大根堆小根堆 Top-K问题 Java实现

关键词由CSDN通过智能技术生成

500家公司求出前100强

典型的top-k问题可以通过建立堆（优先级队列）来实现，在500家公司求出前100强的问题当中，可以建立一个大根堆，这样你可以通过弹出堆顶的元素存入数组来记录前一百强公司，因为大根堆堆顶元素为最大值，每弹出一次，都会再次调整为大根堆，这样就可以保证每次弹出都是目前堆中最大的元素。（不是很好的做法，下面还有更好的做法）

代码如下：

import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;
public class Test {
    /**
     * array为500家公司
     * store为待存入的前100强公司
     * k为求入前几强
     */
    public static void top10(int[] array, int[]store, int k){
        PriorityQueue<Integer> priorityQueue = new PriorityQueue<Integer>(new Comparator<Integer>() {
            @Override
            public int compare (Integer o1, Integer o2){
                return o2 - o1;
            }
        });
        for(int i = 0; i < array.length; i++){
            priorityQueue.offer(array[i]);
        }
        for(int i = 0; i < k; i++){
            store[i] = priorityQueue.poll();

        }
    }
    public static void main(String[] args){
        /**
         * 测试用例
         * 假设有10家公司，求前5家
         */
        int[] array = {3,1,2,9,6,7,8,4,5,10};
        int[] store = new int[5];
        top10(array, store, 5);
        System.out.println(Arrays.toString(store));
    }
}

实际上这样写还不是很好的办法，当公司的数量特别大时，每进行一次调整都可能浪费很多时间，所以，其实可以建立一个大小为100的小根堆，先将数组的前100个元素放入，再将后面的元素与堆顶元素比较，若比堆顶元素大，则存入，最后剩下堆中的元素就是前100强公司。

代码如下：

import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;
public class Test {
    /**
     * array为500家公司
     * k为前100强的
     */
    public static int[] top10(int[] array, int k){
        int[] ret = new int[k];
        PriorityQueue<Integer> priorityQueue = new PriorityQueue<Integer>(new Comparator<Integer>() {
            @Override
            public int compare (Integer o1, Integer o2){
                return o1 - o2;
            }
        });
        for(int i = 0; i < array.length; i++){
            if(priorityQueue.size() < k){
                priorityQueue.offer(array[i]);
            }else{
                if(priorityQueue.peek() < array[i]){
                    priorityQueue.poll();
                    priorityQueue.offer(array[i]);
                }
            }
        }
        for(int i = 0; i < k ; i++){
            ret[i] = priorityQueue.poll();
        }
        return ret;
    }
    public static void main(String[] args){
        /**
         * 测试用例
         * 假设有10家公司，求前三家强的公司
         */
        int[] array = {3,1,2,9,6,7,8,4,5,10};
        int[] store = top10(array, 3);
        System.out.println(Arrays.toString(store));
    }
}

500家公司里最后100家公司

由上例不难想出，只需要将大根堆修改成小根堆，比较值修改即可

如下代码：

import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;
public class Test {
    /**
     * array为500家公司
     * k为后100弱的
     */
    public static int[] top10(int[] array, int k){
        int[] ret = new int[k];
        PriorityQueue<Integer> priorityQueue = new PriorityQueue<Integer>(new Comparator<Integer>() {
            @Override
            public int compare (Integer o1, Integer o2){
                return o2 - o1;
            }
        });
        for(int i = 0; i < array.length; i++){
            if(priorityQueue.size() < k){
                priorityQueue.offer(array[i]);
            }else{
                if(priorityQueue.peek() > array[i]){
                    priorityQueue.poll();
                    priorityQueue.offer(array[i]);
                }
            }
        }
        for(int i = 0; i < k ; i++){
            ret[i] = priorityQueue.poll();
        }
        return ret;
    }
    public static void main(String[] args){
        /**
         * 测试用例
         * 假设有10家公司，求后三家公司
         */
        int[] array = {3,1,2,9,6,7,8,4,5,10};
        int[] store = top10(array, 3);
        System.out.println(Arrays.toString(store));
    }
}

前五百家公司，第一百强的公司

不可以直接排序，如何用堆的思想实现呢？可能你是想建立一个大根堆，然后弹出前99家公司，然后，再弹出一家便是你的所求吗？想象一下，每弹出一次，都有可能会进行向下调整，时间复杂度立马就上来了；要是求第100强的可以建立一个大小为100的小根堆，这样堆顶元素不就是第一百强了吗？建立了大小为100的小根堆后，再将数组的一百个元素放入内，再将数组剩下的元素与堆顶元素进行比较，如果比他大则交换。

代码如下：

import java.util.Comparator;
import java.util.PriorityQueue;
public class Test {
    /**
     * array为500家公司
     * k为第几强的公司
     */
    public static int top10(int[] array, int k){
        PriorityQueue<Integer> priorityQueue = new PriorityQueue<Integer>(new Comparator<Integer>() {
            @Override
            public int compare (Integer o1, Integer o2){
                return o1 - o2;
            }
        });
        for(int i = 0; i < array.length; i++){
            if(priorityQueue.size() < k){
                priorityQueue.offer(array[i]);
            }else{
                if(priorityQueue.peek() < array[i]){
                    priorityQueue.poll();
                    priorityQueue.offer(array[i]);
                }
            }
        }
        return priorityQueue.peek();
    }
    public static void main(String[] args){
        /**
         * 测试用例
         * 假设有10家公司，求第3强
         */
        int[] array = {3,1,2,9,6,7,8,4,5,10};
        int top3 = top10(array, 3);
        System.out.println(top3);
    }
}

面试题

和上例中前强500家公司位于后面的100家公司是一个道理，但要注意的是小心此题k = 0的情况。

如下代码：

class Solution {
    public int[] smallestK(int[] arr, int k) {
        int[] ret = new int[k];
        if(k == 0){
            return ret;
        }
        PriorityQueue<Integer> priorityQueue = new PriorityQueue<Integer>(new Comparator<Integer>() {
            @Override
            public int compare (Integer o1, Integer o2){
                return o2 - o1;
            }
        });
        for(int i = 0; i < arr.length; i++){
            if(priorityQueue.size() < k){
                priorityQueue.offer(arr[i]);
            }else{
                if(priorityQueue.peek() > arr[i]){
                    priorityQueue.poll();
                    priorityQueue.offer(arr[i]);
                }
            }
        }
        for(int i = 0; i < k ; i++){
            ret[i] = priorityQueue.poll();
        }
        return ret;
    }
}