4种解法秒杀TopK(快排变形/堆/二叉搜索树/计数排序)

一、原题描述:

二、题目分析:

TopK问题,不管是求前K大/前K小/第K大/第K小等,都有4种不错的方法喔:

1. O(N):用快排变形最最最高效解决TopK问题

2. O(NlogK):大根堆(前K小)/小根堆(前K大)

3. O(NlogK):二叉搜索树

4. O(N): 对于数据范围有限的情况例如本题,可以直接计数排序O(N)高效解决~

下面针对本题,求数组中的前K小的数,用上面的4种方法分别实现一遍。

1. 用快排变形最最最高效解决TopK问题 O(N)

注意找前K大/前K小/第K大/第K小,是不需要对整个数组进行O(NlogN)的排序的!因为可以通过快排切分直接O(N)找到第K大的数(比如求中位数就可以用本方法O(N)找到第mid大的数,如果只会先排序再找的话,那啥...基本上就交代了叭( ͡° ͜ʖ ͡°)...)。

因此本题先通过快排切分排好第K小的数,根据快排切分的性质,它左边的K - 1个数都小于等于它,因此它以及它左边的数就是我们要找的前K小的数。

下面代码给出了详细的注释,没啥好啰嗦的,就是快排模版要记牢哈~

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }
        // ⚠️注意最后一个参数传入我们要找的下标(第k小的数下标是k-1)
        return quickSearch(arr, 0, arr.length - 1, k - 1);
    }

    private int[] quickSearch(int[] nums, int lo, int hi, int k) {
        // 每快排切分1次,找到排序后下标为j的元素,如果j恰好等于k就返回j以及j左边所有的数;
        int j = partition(nums, lo, hi);
        if (j == k) {
            return Arrays.copyOf(nums, j + 1);
        }
        // 否则根据下标j与k的大小关系来决定继续切分左段还是右段。
        return j > k? quickSearch(nums, lo, j - 1, k): quickSearch(nums, j + 1, hi, k);
    }

    // 快排切分,返回下标j,使得比nums[j]小的数都在j的左边,比nums[j]大的数都在j的右边。
    private int partition(int[] nums, int lo, int hi) {
        int v = nums[lo];
        int i = lo, j = hi + 1;
        while (true) {
            while (++i <= hi && nums[i] < v);
            while (--j >= lo && nums[j] > v);
            if (i >= j) {
                break;
            }
            int t = nums[j];
            nums[j] = nums[i];
            nums[i] = t;
        }
        nums[lo] = nums[j];
        nums[j] = v;
        return j;
    }
}

快排变形时间复杂度分析:
因为我们是要找下标为k的元素,第一次切分的时候需要遍历整个数组(0 ~ n)找到了下标是j的元素,假如k比j小的话,那么我们下次切分只要遍历数组(0~k-1)的元素就行啦,反之如果k比j大的话,那下次切分只要遍历数组(k+1~n)的元素就行啦,总之平均情况下,可以看作每次调用partition遍历的元素数目都是上一次遍历的1/2,因此时间复杂度是N + N/2 + N/4 + ... + N/N = 2N, 因此时间复杂度是O(N)

2. 大根堆(前K小)/小根堆(前K大) O(NlogK)

用堆虽然时间复杂度比快排变形慢了点,但是因为Java中提供了现成的PriorityQueue(默认小根堆),所以不需要自己写大段的模版代码,因此实现起来最简单,没几行代码,写起来很快~~面试的时候可以先快点写出这个方案

注意本题是求前K小,因此用一个容量为K的大根堆(每次poll出最大的数,那堆中保留的就是前K小啦)。注意不是小根堆嗷!小根堆的话需要把全部的元素都入堆,那是O(NlogN) ,就不是O(NlogK)啦~~

// 保持堆的大小为K,然后遍历数组中的数字,遍历的时候做如下判断:
// 1. 若目前堆的大小小于K,将当前数字放入堆中。
// 2. 否则判断当前数字与大根堆堆顶元素的大小关系,如果当前数字比大根堆堆顶还大(或等于),这个数就直接跳过;
//    反之如果当前数字比大根堆堆顶小,先poll掉堆顶,再将该数字放入堆中。
class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }
        // 默认是小根堆,实现大根堆需要重写一下比较器。
        Queue<Integer> pq = new PriorityQueue<>((v1, v2) -> v2 - v1);
        for (int num: arr) {
            if (pq.size() < k) {
                pq.offer(num);
            } else if (num < pq.peek()) {
                pq.poll();
                pq.offer(num);
            }
        }
        
        // 返回堆中的元素
        int[] res = new int[pq.size()];
        int idx = 0;
        for(int num: pq) {
            res[idx++] = num;
        }
        return res;
    }
}

3. 二叉搜索树也可以O(NlogK)解决TopK问题嗷~

BST相对没有前两种方法辣么热门,但是也很简单,和大根堆的思路差不多~不得不提的是,与前两种方法相比,BST的优势就是求得的前K个数字保证是有序的

因为有重复的数字,所以用的是TreeMap而不是TreeSet(有的语言的标准库自带TreeMultiset,也是可以的)。TreeMap的key是数字,value是该数字的个数。我们遍历数组中的数字,维护一个数字总个数为K的TreeMap,每遍历一个元素:

1. 若目前map中数字个数小于K,则将map中当前数字对应的个数+1;

2. 否则,判断当前数字与map中最大的数字的大小关系:若当前数字大于等于map中的最大数字,就直接跳过该数字;若当前数字小于map中的最大数字,则将map中当前数字对应的个数+1,并将map中最大数字对应的个数减1.

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }
        // TreeMap的key是数字, value是该数字的个数。
        // cnt表示当前map总共存了多少个数字。
        TreeMap<Integer, Integer> map = new TreeMap<>();
        int cnt = 0;
        for (int num: arr) {
            // 1. 遍历数组,若当前map中的数字个数小于k,则map中当前数字对应个数+1
            if (cnt < k) {
                map.put(num, map.getOrDefault(num, 0) + 1);
                cnt++;
                continue;
            } 
            // 2. 否则,取出map中最大的Key(即最大的数字), 判断当前数字与map中最大数字的大小关系:
            //    若当前数字比map中最大的数字还大(或等于),就直接忽略;
            //    若当前数字比map中最大的数字小,则将当前数字加入map中,并将map中的最大数字的个数-1。
            Map.Entry<Integer, Integer> entry = map.lastEntry();
            if (entry.getKey() > num) {
                map.put(num, map.getOrDefault(num, 0) + 1);
                if (entry.getValue() == 1) {
                    map.pollLastEntry();
                } else {
                    map.put(entry.getKey(), entry.getValue() - 1);
                }
            }
            
        }

        // 最后返回map中的元素
        int[] res = new int[k];
        int idx = 0;
        for (Map.Entry<Integer, Integer> entry: map.entrySet()) {
            int freq = entry.getValue();
            while (freq-- > 0) {
                res[idx++] = entry.getKey();
            }
        }
        return res;
    }
}

说到这里了,就额外再说一下吧,Java的标准库中虽然没有TreeMultiset,但是Guava提供了TreeMultiset,因此也可以通过引入Guava包来实现嗷,不是考察重点,简单贴下大概代码。

// ❌leetcode不支持引入第三方包哦
public int[] getLeastNumbers(int[] arr, int k) {
    if (k == 0 || arr.length == 0) {
        return new int[0];
    }
    // TreeMultiset 中允许有重复元素,所以就不需要用TreeMap了。
    TreeMultiset<Integer> set = TreeMultiset.create();
    for(int num: arr) {
        // 1. 遍历每个数字,如果set中的数量小于K,则直接将当前数字加入set中。
        if (set.size() < k) {
            set.add(num);
            continue;
        }
        // 2. 否则判断当前数字与set中最大数字的大小关系:
        //    若当前数字大于等于set中的最大数字,则直接跳过该数字;
        //    若当前数字小于set中的最大数字,则将当前数字加入set,并将set中最大数字的个数-1。
        Multiset.Entry<Integer> lastEntry = set.lastEntry();
        if (num < lastEntry.getElement()) {
            set.remove(lastEntry.getElement(), 1);
            set.add(num);
        }
    }

    // 返回set中的元素
    int[] res = new int[k];
    int idx = 0;
    for(int num: set) {
        res[idx++] = num;
    }
    return res;
}

4. 对于数据范围有限的情况例如本题,可以直接计数排序O(N)高效解决~

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }
        // 统计每个数字出现的次数
        int[] counter = new int[10001];
        for (int num: arr) {
            counter[num]++;
        }
        // 根据counter数组从头找出k个数作为返回结果
        int[] res = new int[k];
        int idx = 0;
        for (int num = 0; num < counter.length; num++) {
            while (counter[num]-- > 0 && idx < k) {
                res[idx++] = num;
            }
            if (idx == k) {
                break;
            }
        }
        return res;
    }
}

三、类似题目:

215.数组中的第K个最大元素

  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值