算法题数据流中的第 K 大元素-深圳市維司達科技有限公司

数据流中的第 K 大元素

问题描述

设计一个找到数据流中第k大元素的类（class）。注意，这是指在已排序的顺序中处于第k个位置的元素，而不是第k个不同的元素。

请实现KthLargest类：

KthLargest(int k, int[] nums)使用整数k和整数流nums初始化对象。
int add(int val)将val插入数据流nums后，返回当前数据流中第k大的元素。

示例：

输入: ["KthLargest", "add", "add", "add", "add", "add"] [[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]] 输出: [null, 4, 5, 5, 8, 8] 解释: KthLargest kthLargest = new KthLargest(3, [4, 5, 8, 2]); kthLargest.add(3); // return 4 (数据流: [2,3,4,5,8]) kthLargest.add(5); // return 5 (数据流: [2,3,4,5,5,8]) kthLargest.add(10); // return 5 (数据流: [2,3,4,5,5,8,10]) kthLargest.add(9); // return 8 (数据流: [2,3,4,5,5,8,9,10]) kthLargest.add(4); // return 8 (数据流: [2,3,4,4,5,5,8,9,10])

算法思路

最小堆：

核心思想：

维护一个大小为k的最小堆，堆顶元素就是第k大的元素
堆中始终保存数据流中最大的k个元素
当堆大小超过k时，移除堆顶（最小的元素）

为什么使用最小堆？

最小堆的堆顶是堆中最小的元素
如果堆中有k个元素，堆顶就是第k大的元素
当有新元素加入时，如果新元素比堆顶大，它会替换掉堆顶，保持堆中始终是最大的k个元素

代码实现

方法一：最小堆

importjava.util.PriorityQueue;classKthLargest{privateintk;privatePriorityQueue<Integer>minHeap;/** * 初始化KthLargest对象 * * @param k 第k大的元素 * @param nums 初始数据流 */publicKthLargest(intk,int[]nums){this.k=k;// 创建最小堆（PriorityQueue默认是最小堆）this.minHeap=newPriorityQueue<>();// 将初始数组中的元素逐个加入堆for(intnum:nums){add(num);}}/** * 添加元素到数据流并返回第k大元素 * * @param val 要添加的元素 * @return 当前数据流中第k大的元素 */publicintadd(intval){// 将新元素加入堆minHeap.offer(val);// 如果堆大小超过k，移除堆顶（最小元素）if(minHeap.size()>k){minHeap.poll();}// 堆顶就是第k大的元素returnminHeap.peek();}}

算法分析

时间复杂度：
- 初始化：O(N log k)，其中N是初始数组长度
- add操作：O(log k)，堆操作的时间复杂度
空间复杂度：O(k)，堆中最多存储k个元素

算法过程

k=3, nums=[4,5,8,2]

初始化：

添加4：堆=[4]，大小=1≤3
添加5：堆=[4,5]，大小=2≤3
添加8：堆=[4,5,8]，大小=3≤3
添加2：堆=[4,5,8,2] → 大小>3 → 移除2 → 堆=[4,5,8]

堆状态：[4,5,8]（最小堆，堆顶=4）

add操作：

add(3)：
- 堆=[4,5,8,3] → 大小>3 → 移除3 → 堆=[4,5,8]
- 返回堆顶=4
add(5)：
- 堆=[4,5,8,5] → 大小>3 → 移除4 → 堆=[5,5,8]
- 返回堆顶=5
add(10)：
- 堆=[5,5,8,10] → 大小>3 → 移除5 → 堆=[5,8,10]
- 返回堆顶=5
add(9)：
- 堆=[5,8,10,9] → 大小>3 → 移除5 → 堆=[8,9,10]
- 返回堆顶=8
add(4)：
- 堆=[8,9,10,4] → 大小>3 → 移除4 → 堆=[8,9,10]
- 返回堆顶=8

最终数据流：[2,3,4,4,5,5,8,9,10]
最大的3个元素：[8,9,10]
第3大元素：8

测试用例

publicclassTestKthLargest{publicstaticvoidmain(String[]args){// 测试用例1：标准示例KthLargestkthLargest1=newKthLargest(3,newint[]{4,5,8,2});System.out.println("Test 1:");System.out.println("add(3): "+kthLargest1.add(3));// 4System.out.println("add(5): "+kthLargest1.add(5));// 5System.out.println("add(10): "+kthLargest1.add(10));// 5System.out.println("add(9): "+kthLargest1.add(9));// 8System.out.println("add(4): "+kthLargest1.add(4));// 8System.out.println();// 测试用例2：k=1（找最大值）KthLargestkthLargest2=newKthLargest(1,newint[]{});System.out.println("Test 2 (k=1):");System.out.println("add(1): "+kthLargest2.add(1));// 1System.out.println("add(3): "+kthLargest2.add(3));// 3System.out.println("add(2): "+kthLargest2.add(2));// 3System.out.println("add(4): "+kthLargest2.add(4));// 4System.out.println();// 测试用例3：k等于数组长度KthLargestkthLargest3=newKthLargest(2,newint[]{0});System.out.println("Test 3 (k=2):");System.out.println("add(-1): "+kthLargest3.add(-1));// -1System.out.println("add(1): "+kthLargest3.add(1));// 0System.out.println("add(-2): "+kthLargest3.add(-2));// -1System.out.println("add(-4): "+kthLargest3.add(-4));// -1System.out.println("add(3): "+kthLargest3.add(3));// 0System.out.println();// 测试用例4：大量重复元素KthLargestkthLargest4=newKthLargest(3,newint[]{5,5,5});System.out.println("Test 4:");System.out.println("add(5): "+kthLargest4.add(5));// 5System.out.println("add(5): "+kthLargest4.add(5));// 5System.out.println("add(10): "+kthLargest4.add(10));// 5System.out.println();// 测试用例5：负数KthLargestkthLargest5=newKthLargest(2,newint[]{-1,-2,-3});System.out.println("Test 5:");System.out.println("add(-4): "+kthLargest5.add(-4));// -2System.out.println("add(0): "+kthLargest5.add(0));// -1System.out.println("add(1): "+kthLargest5.add(1));// 0System.out.println();// 测试用例6：空初始数组KthLargestkthLargest6=newKthLargest(2,newint[]{});System.out.println("Test 6:");System.out.println("add(1): "+kthLargest6.add(1));// 1 (堆大小<2)System.out.println("add(2): "+kthLargest6.add(2));// 1 (堆=[1,2])System.out.println("add(3): "+kthLargest6.add(3));// 2 (堆=[2,3])System.out.println("add(0): "+kthLargest6.add(0));// 2 (堆=[2,3])}}