topK问题

100 阅读2分钟

本文已参与「新人创作礼」活动,一起开启掘金创作之路。

问题描述

就是我们有很多的数据需要整理,这个时候,我们需要选出最大K个数,或者最小K个数,这个时候怎么办。

思路

这个就要想到我们堆的思想,我们的堆是把最大或者最小的放在最上面,我们可以先把所有数据里最前面的K个数放到我们的堆里面,后面的数据跟堆最上面的数比较,看看是否符合条件,如果符合条件,就把第一个丢掉,符合条件的元素放过来。

案例

数组【3,8,0,9,2,5,6,1,4,7】

这次我们要挑选最小的三个数

我们就先把前三个数放到小根堆里面。

屏幕截图 2022-10-01 180630.jpg

这个时候我们挑选9,来看看能不能入堆,那我们入堆的条件是什么呢?

因为我们是挑选最小的三个,所以我们是小于最前面的就可以入堆了。

但是我们9明显不符合,所以不能入堆,

这个时候我们看到了2,我们发现好像是小于最上面的那个值的,所以我们就可以让它入堆了。

屏幕截图 2022-10-01 180841.jpg

然后我们让后面的数也经历判断是否入堆,如果入堆,就把堆重排,最后留在堆里面的就是最小的三个

代码

public static int[] topK(int[] arr , int k){
    if(arr == null || arr.length < k){
        return arr;
    }
    int[] ans = new int[k];
    MyBigHeap myBigHeap = new MyBigHeap(k);
    for(int i = 0;i<k;i++){
        myBigHeap.push(arr[i]);
    }
    for(int i = k;i<arr.length;i++){
        int com = myBigHeap.pop();
        if(arr[i] < com){
            myBigHeap.push(arr[i]);
        }else{
            myBigHeap.push(com);
        }
    }
    for(int i = 0;i<k;i++){
        ans[i] = myBigHeap.pop();
    }
    return ans;
}

这上面就是topK问题的解答,至于这边的MyBigHeap

大家可以去大堆和小堆的实现 - 掘金 (juejin.cn)