703. Kth Largest Element in a Stream

──────────────────────────────────────────────────────────────────────

設計一個類別，add(val) 每次加入一個數字並回傳目前為止第 k 大的值。

思路

這是 Top-K 的資料流版本：資料不斷進來，每次都要立刻回答「第 k 大是誰」。

最直接的想法是每次 add 都排序一遍，但那是每次 $O (n \log n)$ 。要更快就得問：我真的需要記住全部的數字嗎？

不需要。答案永遠只跟目前最大的那 k 個有關 —— 第 k+1 大以下的數字永遠不可能變成答案（後面只會有更多數字進來，排名只會往後掉）。所以：

維持一個大小恰好為 k 的最小堆，堆頂就是答案。

用最小堆而不是最大堆是這題的關鍵顛倒：堆裡裝的是「前 k 名」，而堆頂是這 k 個裡最小的，也就是第 k 大 —— 剛好是要回傳的東西。同時它也是「門檻」：新來的數字比它大就該進來、把它擠掉；比它小就直接丟掉。

add 的動作因此只有兩步：先推進去，超過 k 個就彈掉最小的。堆的大小自動維持在 k。

這個「大小為 k 的最小堆」就是 347. Top K Frequent Elements 用的同一招，只是那題資料是一次給完的，這題是流進來的 —— 而堆的好處正是它不需要看到全部資料就能維持狀態。

解題方向

class KthLargest:

    def __init__(self, k: int, nums: List[int]):
        self.k = k
        self.heap = []
        for num in nums:
            self.add(num)

    def add(self, val: int) -> int:
        heapq.heappush(self.heap, val)
        while len(self.heap) > self.k:
            heapq.heappop(self.heap)
        
        return self.heap[0]

建構子直接重複呼叫 add 來建立初始狀態，不需要另外寫一套邏輯 —— 這是最乾淨的版本，因為「維持大小 k」的規則只寫在一個地方。

另外兩種等價的寫法：

    # 版本 A：建構子先 heapify，靠 add 慢慢修剪
    def __init__(self, k, nums):
        self.k = k
        self.nums = nums
        heapq.heapify(self.nums)          # O(n)，但沒有立刻修剪

    # 版本 B：建構子 heapify 之後就修剪到 k
    def __init__(self, k, nums):
        self.k = k
        self.nums = nums
        heapq.heapify(self.nums)
        while len(self.nums) > k:
            heapq.heappop(self.nums)

三種的結果完全相同（我都驗證過）。差別在初始化的成本和記憶體：

版本 A 的堆一開始有 n 個元素，要到第一次 add 才被修剪 —— 如果 n 很大而 add 呼叫得少，中間會一直佔著 $O (n)$ 記憶體
版本 B 立刻修剪到 k，之後穩定佔 $O (k)$
上面主要的那版每次只推一個就修剪，堆從頭到尾不超過 k + 1

三者的初始化時間同級（都是 $O (n \log k)$ 或 $O (n)$ heapify 加修剪），但版本 B 和主要那版的記憶體上限是 $O (k)$ ，版本 A 是 $O (n)$ 。面試時選一個並說得出理由就好。

while 還是 if？ 每次只推一個元素，所以最多只會超出一個，用 if 就夠。寫 while 更保險（例如建構子那種一次塞很多的情況），兩者都對。

補充

self.heap[0] 為什麼是第 k 大？ 因為堆裡恰好是「目前最大的 k 個」，而最小堆的堆頂是其中最小的 —— 「前 k 大裡最小的那個」就是第 k 大。

題目保證呼叫 add 時總元素數至少有 k 個，所以不用處理「堆還不滿 k 個」的情況。如果沒有這個保證，就要在 add 裡多判斷 len(self.heap) < self.k 時回傳 None 或拋錯。

同一組工具的題目：347. Top K Frequent Elements（一次性的 Top-K）、215. Kth Largest Element in an Array（靜態陣列版，也可以用 quickselect 做到平均 $O (n)$ ）、295. Find Median from Data Stream（資料流求中位數，改用兩個堆互相平衡）。

295 特別值得一起看：它和這題都是「資料流 + 堆」，但因為要的是中位數而不是第 k 大，所以需要一個最大堆管左半、一個最小堆管右半。整理見 Heap / Top K 模板。

複雜度

__init__ 時間 $O (n \log k)$ — n 個初始元素各推入一次，每次 $O (\log k)$
add 時間 $O (\log k)$ — 一次 push、最多一次 pop
空間 $O (k)$ — 堆裡永遠只有 k 個元素

其中 n 是初始陣列長度、k 是題目給的名次。

空間是這個解法最漂亮的地方：不管資料流進來多少數字，記憶體都固定在 $O (k)$ —— 這正是「資料流」類題目要的性質，也是為什麼不能用「存起來再排序」。