算法题 数据流中的第 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])

算法思路

最小堆：

核心思想：

1. 维护一个大小为k的最小堆，堆顶元素就是第k大的元素
2. 堆中始终保存数据流中最大的k个元素
3. 当堆大小超过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]

初始化：

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

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

add操作：

1. add(3)：

   • 堆=[4,5,8,3] → 大小>3 → 移除3 → 堆=[4,5,8]
   • 返回堆顶=4

2. add(5)：

   • 堆=[4,5,8,5] → 大小>3 → 移除4 → 堆=[5,5,8]
   • 返回堆顶=5

3. add(10)：

   • 堆=[5,5,8,10] → 大小>3 → 移除5 → 堆=[5,8,10]
   • 返回堆顶=5

4. add(9)：

   • 堆=[5,8,10,9] → 大小>3 → 移除5 → 堆=[8,9,10]
   • 返回堆顶=8

5. 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])}}

关键点

1. 堆的大小：

   • 始终维护堆大小为k
   • 超过k时移除最小元素（堆顶）
   • 确保堆中始终是最大的k个元素

2. 为什么是最小堆而不是最大堆？：

   • 最小堆能快速获取k个最大元素中的最小值（即第k大）
   • 最大堆需要存储所有元素，空间复杂度为O(N)

3. 边界情况处理：

   • 初始数组为空
   • k=1（找最大值）
   • k大于初始数组长度
   • 重复元素

常见问题

1. 为什么不用最大堆？

   • 最大堆需要存储所有元素才能找到第k大，空间复杂度O(N)
   • 最小堆只需存储k个元素，空间效率更高

2. 时间复杂度O(log k)？

   • 堆的插入和删除操作时间复杂度都是O(log size)
   • 堆的大小始终≤k，所以是O(log k)

3. 找第k小元素？

   • 使用最大堆维护最小的k个元素
   • 堆顶就是第k小的元素