Kth Largest trong data stream: tại sao min-heap kích thước k là ranh giới đúng

Bài toán khoác chiếc áo kịch bản tuyển sinh đại học, nhưng bóc tách ra thì chỉ có một câu hỏi lặp đi lặp lại: sau mỗi lần chèn vào một dãy số nguyên đang lớn dần, giá trị lớn thứ k hiện tại là bao nhiêu? Câu hỏi đó phải trả lời được nhanh — tối đa 10^4 lần gọi add — và dữ liệu cứ đến liên tục, nên không thể trì hoãn công việc đến cuối.

Constraint quan trọng nhất là thứ đặt ra mục tiêu: 1 <= k <= nums.length + 1. Bài toán đảm bảo rằng khi bạn được yêu cầu trả về phần tử lớn thứ k, luôn có ít nhất k phần tử trong stream. Tức là bạn không bao giờ phải xử lý tình huống câu trả lời không tồn tại.

Đề bài

Thiết kế class KthLargest nhận một số nguyên k và mảng ban đầu nums, theo dõi stream các giá trị liên tục và trả về phần tử lớn thứ k sau mỗi lần thêm mới. Thứ hạng là theo vị trí — các giá trị trùng lặp được tính là các entry riêng biệt. Full statement: LeetCode 703.

Ràng buộc:
- 0 <= nums.length <= 10^4
- 1 <= k <= nums.length + 1
- -10^4 <= nums[i] <= 10^4
- -10^4 <= val <= 10^4
- Có nhiều nhất 10^4 lần gọi add.

Constraints định hình thiết kế trước khi bạn chọn data structure

Bắt đầu từ nums.length <= 10^4 và tối đa 10^4 lần gọi add. Nếu mỗi add là O(n log n), tổng lượng công việc qua tất cả các lần gọi vào khoảng 10^8 phép toán — con số đó nằm ở vùng ranh giới và gần chắc chắn sẽ TLE trên judge của LeetCode. Bạn cần mỗi add nhiều nhất là O(log n). Thực ra, vì chỉ cần quan tâm đến k phần tử, O(log k) còn tốt hơn.

Dải giá trị [-10^4, 10^4] nhỏ và có dấu. Không lo overflow. Không cần trick với sentinel. Phép so sánh thẳng thắn.

k có thể bằng 1, nghĩa là "luôn trả về maximum đang chạy." Nó cũng có thể bằng nums.length + 1 sau lần add đầu tiên, nghĩa là "trả về minimum của tất cả phần tử cho đến nay." Data structure cần xử lý cả hai cực mà không cần special-case.

Stream có thể chứa duplicate. Điểm số 7 xuất hiện bốn lần là bốn entry riêng biệt. Phần tử lớn thứ k là khái niệm theo vị trí, không phải theo sự phân biệt.

Approach 1: sắp xếp mỗi lần add

Cách đọc đơn giản nhất của bài toán: giữ tất cả phần tử, sắp xếp chúng, trả về vị trí k từ bên phải. Constructor O(1) (chỉ lưu k và list) và add O(n log n).

from typing import List
 
class KthLargest:
    def __init__(self, k: int, nums: List[int]):
        self.k = k
        self.nums = list(nums)
 
    def add(self, val: int) -> int:
        self.nums.append(val)
        self.nums.sort(reverse=True)
        return self.nums[self.k - 1]

using System.Collections.Generic;
 
public class KthLargest {
    private int k;
    private List<int> nums;
 
    public KthLargest(int k, int[] nums) {
        this.k = k;
        this.nums = new List<int>(nums);
    }
 
    public int Add(int val) {
        nums.Add(val);
        nums.Sort((a, b) => b.CompareTo(a));
        return nums[k - 1];
    }
}

Nó đúng về mặt logic. Và cực kỳ chậm với các constraint đã cho: sắp xếp một list n phần tử mỗi lần gọi có nghĩa là tổng công việc là O(total*n * n _ log n), với 10^4 lần add cùng 10^4 phần tử ban đầu thì vào khoảng 10^8 * log(10^4), vượt xa giới hạn chấp nhận được.

Vấn đề sâu hơn là việc sắp xếp bỏ đi và tính lại thông tin. Sau lần sort đầu tiên, bạn đã biết thứ tự; lần gọi tiếp theo thêm một phần tử, và bạn sort lại từ đầu. Bạn đang lãng phí công việc đã làm.

Time: O(n log n) mỗi add, O(1) cho constructor — nhưng thực tế bạn sẽ pre-sort trong constructor nếu muốn tối ưu vài lần gọi đầu, làm constructor cũng thành O(n log n).

Space: O(n) để lưu tất cả phần tử, tăng không giới hạn khi add được gọi.

Approach 2: min-heap kích thước k — ranh giới đúng

Insight là bạn không cần thứ tự sắp xếp đầy đủ. Bạn cần đúng một con số: phần tử lớn thứ k. Và phần tử lớn thứ k chính là phần tử nhỏ nhất trong số k phần tử lớn nhất.

Một min-heap kích thước chính xác k lưu k phần tử lớn nhất bạn đã thấy. Root của nó — minimum trong số k phần tử đó — theo định nghĩa là phần tử lớn thứ k tổng thể. Mỗi khi có phần tử mới đến, một trong hai điều xảy ra: nó nhỏ hơn phần tử lớn thứ k hiện tại (không thể đẩy bất kỳ phần tử nào trong top k ra, nên loại bỏ nó) hoặc nó lớn hơn hoặc bằng (nó vào top k, và phần tử nhỏ nhất của top k bị trục xuất). Heap thực hiện cả hai bước trong O(log k).

import heapq
from typing import List
 
class KthLargest:
    def __init__(self, k: int, nums: List[int]):
        self.k = k
        self.heap = list(nums)
        heapq.heapify(self.heap)          # O(n) linear heapify
        while len(self.heap) > k:
            heapq.heappop(self.heap)      # trục xuất nhỏ nhất cho đến khi size == k
 
    def add(self, val: int) -> int:
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)
        return self.heap[0]               # root là phần tử lớn thứ k

using System.Collections.Generic;
 
public class KthLargest {
    private int k;
    private PriorityQueue<int, int> minHeap;
 
    public KthLargest(int k, int[] nums) {
        this.k = k;
        this.minHeap = new PriorityQueue<int, int>();
        foreach (int num in nums) {
            minHeap.Enqueue(num, num);    // priority == value -> min-heap behavior
        }
        while (minHeap.Count > k) {
            minHeap.Dequeue();
        }
    }
 
    public int Add(int val) {
        minHeap.Enqueue(val, val);
        if (minHeap.Count > k) {
            minHeap.Dequeue();
        }
        return minHeap.Peek();
    }
}

Time: Constructor là O(n) cho heapify cộng O((n - k) log n) cho vòng trục xuất, tổng O(n log n) trong trường hợp xấu nhất. Mỗi add là O(log k) — push và nhiều nhất một pop trên heap kích thước k.

Space: O(k). Heap không bao giờ vượt quá k + 1 entry (bạn push rồi pop ngay nếu vượt).

Trace tay: k = 3, nums = [4, 5, 8, 2]

Đi qua constructor và từng lần gọi add từ Example 1.

Từng bước cho hai lần add đầu tiên:

Trạng thái ban đầu sau constructor: heap = [4, 5, 8] (nội bộ là min-heap; root là 4). Ba phần tử lớn nhất của {4, 5, 8, 2} là 8, 5, 4. Phần tử lớn thứ k (thứ 3) là 4, nằm ở root.

add(3): push 3 vào heap -> [3, 4, 5, 8], size 4 > k=3, pop minimum (3) -> heap trở về [4, 5, 8]. Root = 4. Con số 3 không đủ lớn để vào top 3 nên bị trục xuất ngay. Return 4. Đúng.

add(5): push 5 -> [4, 5, 5, 8], pop minimum (4) -> [5, 5, 8]. Root = 5. Con số 5 mới đến đẩy phần tử lớn thứ 3 cũ (4) ra ngoài. Return 5. Đúng.

add(10): push 10 -> [5, 5, 8, 10], pop minimum (5) -> [5, 8, 10]. Root = 5. Return 5. Con số 10 vào top 3 (10, 8, 5); phần tử lớn thứ k là 5.

add(9): push 9 -> [5, 8, 9, 10], pop 5 -> [8, 9, 10]. Root = 8. Return 8.

add(4): push 4 -> [4, 8, 9, 10], pop 4 -> [8, 9, 10]. Root = 8. Return 8.

Dãy output: [4, 5, 5, 8, 8]. Khớp.

Các edge case cần lý luận rõ ràng

nums rỗng khi construction: heapify([]) tạo ra heap rỗng. Vòng while len(heap) > k sai ngay từ đầu. Vài lần add đầu tiên xây dựng heap từ đầu. Đảm bảo 1 <= k <= nums.length + 1 cho thấy sau đủ lần add, heap sẽ đạt kích thước k trước khi bài toán yêu cầu bất kỳ thứ gì.

k = 1: Heap lưu đúng một phần tử: maximum đang chạy. Mỗi add lớn hơn maximum hiện tại sẽ trục xuất nó; các giá trị nhỏ hơn bị loại bỏ ngay. Root của heap luôn là maximum.

k bằng tổng số phần tử: Khi k bằng số phần tử bạn cuối cùng giữ, heap lưu tất cả, và root là global minimum. Điều này làm giảm hiệu quả về mặt space, nhưng time mỗi add vẫn là O(log k) = O(log n).

Duplicate: [7, 7, 7, 7, 8, 3] với k = 4 từ Example 2. Sau constructor, heap giữ 4 phần tử lớn nhất: [7, 7, 7, 8], root = 7. Thêm 2: 2 < 7, bị trục xuất ngay, return 7. Thêm 10: heap thành [7, 7, 7, 8, 10], pop 7 -> [7, 7, 8, 10], return 7. Điều này cho thấy duplicate được tính là các entry riêng biệt — phần tử lớn thứ k theo vị trí là 7 dù có bốn số 7.

Giá trị âm: -10^4 đến 10^4. Các phép so sánh trong heap hoàn toàn là số học; số âm hoạt động giống hệt số dương. Không cần xử lý đặc biệt.

Invariant làm cho approach này hoạt động

Heap duy trì một invariant: nó chứa đúng k phần tử lớn nhất đã thấy cho đến nay. Invariant đó được thiết lập trong constructor (trục xuất cho đến khi size = k) và được bảo toàn mỗi add (push rồi trục xuất nếu vượt). Với invariant đó, root luôn là minimum trong số k phần tử lớn nhất, tức là chính xác phần tử lớn thứ k theo định nghĩa.

Nhận thức quan trọng là bạn không theo dõi danh sách đã sắp xếp — bạn đang theo dõi một ranh giới. Ranh giới đó là: "phần tử này có thuộc top k không?" Root của heap là giá trị ranh giới đó. Bất kỳ thứ gì tại hoặc trên root đều thuộc top k; bất kỳ thứ gì thấp hơn thì không. Mỗi add hoặc cập nhật ranh giới (phần tử mới vào, ranh giới cũ ra) hoặc giữ nguyên nó (phần tử mới thấp hơn ranh giới, bị loại bỏ).

Bảng so sánh

Approach nào tôi sẽ ship, và tại sao

Min-heap. Không cần phải suy nghĩ. Brute force đúng và tôi chấp nhận nó như bước đầu tiên trong phỏng vấn, nhưng tôi sẽ không dừng ở đó khi add được gọi 10^4 lần. Heap giảm công việc mỗi lần gọi từ O(n log n) xuống O(log k), và dùng O(k) space thay vì O(n). Nếu k nhỏ so với n — đây là trường hợp điển hình — cả hai sự khác biệt đều đáng kể.

Tình huống duy nhất brute force có thể tranh luận là nếu bạn có k rất lớn (gần bằng n) và data stream ngắn. Trong trường hợp đó heap không tiết kiệm nhiều bộ nhớ, và code đơn giản hơn có thể dễ maintain hơn. Nhưng các constraint bài toán làm cho k nhỏ trong thực tế, nên heap là lựa chọn đúng.

Trong production bạn sẽ dùng priority queue của ngôn ngữ (Python heapq, C# PriorityQueue, Java PriorityQueue). Tất cả đều implement binary min-heap bên dưới. Pattern — push rồi có điều kiện pop để maintain size k — giống nhau trong mọi ngôn ngữ.

Vị trí trong heap series

Đây là bài thứ năm trong series heap và priority queue, sau Last Stone Weight (1046), K Closest Points to Origin (973), Kth Largest Element in an Array (215), và Task Scheduler (621). Mỗi bài đó giới thiệu một góc độ khác nhau về heap pattern.

Last Stone Weight dùng max-heap để liên tục lấy ra hai phần tử lớn nhất. K Closest Points dùng max-heap kích thước k để theo dõi những điểm gần nhất theo khoảng cách. Kth Largest Element in an Array (215) tìm phần tử lớn thứ k trong mảng tĩnh — cùng trick min-heap-kích-thước-k áp dụng ở đó, làm cho 703 là bước tiếp theo tự nhiên khi mảng trở thành stream trực tiếp. Task Scheduler giới thiệu ý tưởng dùng heap để lên lịch công việc tối ưu.

703 thêm chiều streaming: dữ liệu đến theo từng bước, và bạn cần câu trả lời sau mỗi lần chèn duy nhất. Chi phí O(log k) của heap mỗi operation là thứ làm streaming khả thi.

215. Kth Largest Element in an Array là phiên bản tĩnh của bài toán này. Stream được thay bằng mảng cố định; bạn có thể giải với cùng min-heap kích thước k hoặc quickselect. Approach heap chuyển trực tiếp.

295. Find Median from Data Stream nâng độ khó: thay vì phần tử lớn thứ k, duy trì median. Giải pháp dùng hai heap — max-heap cho nửa dưới và min-heap cho nửa trên — và rebalance sau mỗi lần chèn.

347. Top K Frequent Elements chuyển tiêu chí từ giá trị sang tần suất. Bạn theo dõi k phần tử không phải theo giá trị thô mà theo số lần xuất hiện. Heap kích thước k vẫn áp dụng, nhưng key so sánh thay đổi.

378. Kth Smallest Element in a Sorted Matrix áp dụng câu hỏi "tìm phần tử thứ k" vào cấu trúc 2D nơi các hàng và cột đều được sắp xếp. Min-heap dẫn dắt quá trình duyệt.