TopK 问题

TopK 问题

数组中的第K个最大元素

在未排序的数组中找到第 k 个最大的元素。

示例 :

输入: [3, 2, 1, 5, 6, 4] 和 k = 2
输出: 5

堆：维护一个大小为 K 的最小堆, 最小堆中的元素就是第 K 大元素。

public int findKthLargest(int[] nums, int k) {
    // 一个基于优先级的无界优先级队列, 按照自然规律排序
    PriorityQueue<Integer> heap =  new PriorityQueue<Integer>(); 
    
    for (int n: nums) {
        heap.offer(n);
        if (heap.size() > k)  // 保持第K大元素在队列顶部
            heap.poll();
    }
    return heap.poll();        
}

快速选择：第 K 个元素的问题, 可以使用快速排序实现。

public int findKthLargest(int[] nums, int k) {
    int l = 0;
    int h = nums.length - 1;
    k = nums.length + 1 - k; //第 k 大 == 第 n + 1 - k 小
    while(l < h){
        int pos = partition(nums, l, h);
        if(pos == k - 1){
            break;
        } else if(pos < k - 1) {
            l = pos + 1;
        } else {
            h = pos - 1;
        } 
    }
    return nums[k - 1];
}

private int partition(int[] arr, int left, int right) {
    // 设定基准值（pivot）
    int pivot = arr[left];
    int i = left;
    int j = right;
    while(i < j) {
        // 从右向左找小于x的数来填s[i]
		while(i < j && s[j] >= x) {
            j--;
        }  
		if(i < j) {
			s[i] = s[j]; //将s[j]填到s[i]中, s[j]就形成了一个新的坑
			i++;
		}
 
		// 从左向右找大于或等于x的数来填s[j]
		while(i < j && s[i] < x) {
            i++;
        } 
		if(i < j) {
			s[j] = s[i]; //将s[i]填到s[j]中, s[i]就形成了一个新的坑
			j--;
		}        
    }
    arr[i] = pivot;
    return i;
}

前 K 个高频元素

给定一个非空的整数数组, 返回其中出现频率前 k 高的元素。

示例 :

输入: nums = [1, 1, 1, 2, 2, 3], k = 2
输出: [1, 2]

思路：设置若干个桶, 每个桶存储出现频率相同的数。桶的下标表示数出现的频率, 即第 i 个桶中存储的数出现的频率为 i。

把数都放到桶之后, 从后向前遍历桶, 最先得到的 k 个数就是出现频率最多的的 k 个数

public List<Integer> topKFrequent(int[] nums, int k) {
    //用于计数的map
    HashMap<Integer, Integer> count = new HashMap<Integer, Integer>();

    for(int num : nums) { //初始化每个元素出现的频率
        count.put(num, count.getOrDefault(num, 0) + 1);
    }

    //构建大顶堆
    PriorityQueue<Integer> heap = new PriorityQueue<Integer>((n1, n2) -> count.get(n1) - count.get(n2));

    for (int n : count.keySet()) {
        heap.add(n);
        if (heap.size() > k) {
            heap.poll();
        }
    }

    List<Integer> result = new LinkedList<>();
    while(!heap.isEmpty()) {
        result.add(heap.poll());
    }
    Collections.reverse(result); //反转
    return result;
}