LeetCode 703. 数据流中的第 K 大元素

约 5 分钟
文章信息 9 个分类 6 个标签
分类
LeetCode CodeTop 简单 未解决 5-10 腾讯 字节 虾皮 商汤
标签
树 二叉树 设计 堆（优先队列） 二叉搜索树 数据流

题目描述

✅ 703. 数据流中的第 K 大元素

image-20260929001957446

image-20260929001957447

题意分析

初始数组和之后加入的元素共同组成数据流，每次加入后返回按从大到小排序的第 k 项。相同数值每出现一次都占一个位置，因此需要保留重复元素。

解法：固定容量小根堆

核心思路

[!blue]

维护一个最多存放 k 项的小根堆：数据不足 k 项时全部保留，否则只保留当前最大的 k 项。堆满后，堆顶是保留项中的最小值，恰好处于从大到小的第 k 个位置。

加入新值时，堆未满就直接放入。堆已满时，把堆顶看作进入前 k 项的门槛：新值更大，就删除旧堆顶并加入新值；新值不大于堆顶，则已有 k 项不小于它，保留原堆即可。新值等于堆顶时，保留哪一次出现都不影响排名对应的数值。

这样维护后，堆始终保留所需的最大 k 项。之后只会新增数据，第 k 大的门槛不会降低，因此已经拒绝或淘汰的较小项无需重新考虑。每次只读取堆顶作为答案，保留堆内数据供后续添加继续使用。

解题步骤

  1. 保存 k 并创建空的小根堆，构造时依次复用 add 添加初始数组中的元素。
  2. add 中先检查堆大小：少于 k 项时直接入堆。
  3. 已有 k 项时，只有新值大于堆顶才弹出堆顶并加入新值。
  4. 返回当前堆顶。构造过程中可能还不足 k 项，此时内部调用的返回值会被忽略；题目保证 k <= nums.length + 1，所以外部第一次调用 add 后就已有足够的元素。

代码实现

class KthLargest {
    private final int k;
    private final PriorityQueue<Integer> heap;

    public KthLargest(int k, int[] nums) {
        this.k = k;
        heap = new PriorityQueue<>();

        for (int num : nums) {
            add(num);
        }
    }

    public int add(int val) {
        // 未满先接纳，满后只有超过门槛的值才需替换
        if (heap.size() < k) {
            heap.offer(val);
        } else if (val > heap.peek()) {
            // 堆中只保留当前最大的 k 个数，堆顶就是第 k 大。
            heap.poll();
            heap.offer(val);
        }

        // 只查看门槛，不弹出仍需保留的第 k 大
        return heap.peek();
    }
}

import "container/heap"

type IntHeap struct {
    data []int
}

func (h IntHeap) Len() int {
    return len(h.data)
}

func (h IntHeap) Less(i int, j int) bool {
    return h.data[i] < h.data[j]
}

func (h IntHeap) Swap(i int, j int) {
    h.data[i], h.data[j] = h.data[j], h.data[i]
}

func (h *IntHeap) Push(x any) {
    h.data = append(h.data, x.(int))
}

func (h *IntHeap) Pop() any {
    val := h.data[len(h.data)-1]
    h.data = h.data[:len(h.data)-1]
    return val
}

type KthLargest struct {
    k    int
    heap IntHeap
}

func Constructor(k int, nums []int) KthLargest {
    kth := KthLargest{k: k, heap: IntHeap{}}
    heap.Init(&kth.heap)
    for _, num := range nums {
        kth.Add(num)
    }
    return kth
}

func (this *KthLargest) Add(val int) int {
    // 未满先接纳，满后只有超过门槛的值才需替换
    if this.heap.Len() < this.k {
        heap.Push(&this.heap, val)
    } else if val > this.heap.data[0] {
        // 堆中只保留当前最大的 k 个数，堆顶就是第 k 大。
        heap.Pop(&this.heap)
        heap.Push(&this.heap, val)
    }
    // 只查看门槛，不弹出仍需保留的第 k 大
    return this.heap.data[0]
}

复杂度分析

  • 时间复杂度：构造 $O(n\log(k+1))$，其中 $n$ 为初始数组长度；单次添加均摊 $O(\log(k+1))$，包含底层数组偶发扩容的成本。满堆时若新值未超过堆顶，只需 $O(1)$。
  • 空间复杂度：$O(k)$，堆不超过 k 项。

关键点总结

[!green]

  • 第 k 大转化为最大 k 项中的最小值。
  • 满堆时拒绝较小值不会影响未来结果。
  • k = 1 时同一逻辑只保留最大值；初始数组为空时，题目约束保证 k = 1，第一次加入后即可正常返回。

易错点总结

[!yellow]

  • 使用大根堆返回的会是最大值，而不是保留集合的门槛。
  • 不比较就先弹旧堆顶再加入较小值，会丢失本应保留的数。
  • 查看答案时弹出，破坏后续维护。

相似题目

题目 难度 关联与区别
215. 数组中的第K个最大元素 中等 原题一次性求数组第k大，本题持续加入数据，需要维持大小k的小顶堆。
295. 数据流的中位数 困难 同样动态查询排名边界，中位数位置随总数变化，需要维护两半而不是固定k项。
347. 前 K 个高频元素 中等 用大小受限的堆保留排名靠前的候选；本题支持持续插入时维护第 k 大，该题以元素频次为排序依据。
692. 前K个高频单词 中等 用大小受限的堆保留排名靠前的候选；本题支持持续插入时维护第 k 大，该题以单词频次和字典序联合排序。
973. 最接近原点的 K 个点 中等 用大小受限的堆保留排名靠前的候选；本题支持持续插入时维护第 k 大，该题以到原点的平方距离为排序依据。
转载与许可
作者
链接 https://hgnulb.github.io/blog/2023/51429010
许可 本博客所有文章除特别声明外，均采用 CC BY-NC-SA 4.0 许可协议，转载请注明出处！