K Closest Points to Origin: sort toàn bộ, min heap, và khi nào nên dùng bounded max heap

Công thức khoảng cách ở bài này là sqrt(x² + y²), và bản năng đầu tiên là sort theo nó. Cách đó hoạt động. Nhưng sort là O(n log n) bất kể bạn cần 1 kết quả hay n-1, và với n lớn, k nhỏ thì có một hình dạng giải pháp tốt hơn. Bài này nằm trong series chính xác vì ba approach — sort, min heap, bounded max heap — thể hiện một trade-off ladder thực sự, không chỉ là biến thể về mặt phong cách.

Những gì constraint buộc bạn phải làm

Input là 1 <= k <= points.length <= 10^4, với tọa độ trong khoảng [-10^4, 10^4]. Từ đó suy ra được một vài điều.

Giới hạn tọa độ có nghĩa là x² + y² nằm vừa trong một 32-bit integer: trường hợp xấu nhất là 10^4 * 10^4 + 10^4 * 10^4 = 2 * 10^8, vẫn còn xa Int32.MaxValue (khoảng 2.1 * 10^9). Vì vậy bạn có thể dùng bình phương khoảng cách Euclidean làm khóa so sánh mà không bị overflow và không có nhiễu floating-point từ Math.Sqrt. Tôi bỏ sqrt hoàn toàn bên dưới — nó không ảnh hưởng đến thứ tự, chỉ ảnh hưởng đến độ lớn.

n <= 10^4 nghĩa là O(n log n) đủ nhanh. Không có approach nào trong ba cách bên dưới gây TLE ở đây. Sự khác biệt giữa chúng nằm ở khi n lớn và k nhỏ, đó chính xác là trường hợp max-heap approach được thiết kế cho. Với n = 10^4, sự khác biệt không đáng kể. Với n = 10^7 và k = 10, đó là khoảng cách giữa việc giữ mười triệu item trong bộ nhớ hay chỉ mười.

Approach 1: sort toàn bộ

Sort tất cả điểm theo bình phương khoảng cách, lấy k điểm đầu tiên. Đơn giản, dễ đọc, và đúng.

class Solution:
    def kClosest(self, points: list[list[int]], k: int) -> list[list[int]]:
        points.sort(key=lambda p: p[0] * p[0] + p[1] * p[1])
        return points[:k]

public class Solution {
    public int[][] KClosest(int[][] points, int k) {
        Array.Sort(points, (p1, p2) => {
            int d1 = p1[0] * p1[0] + p1[1] * p1[1];
            int d2 = p2[0] * p2[0] + p2[1] * p2[1];
            return d1.CompareTo(d2);
        });
 
        int[][] result = new int[k][];
        for (int i = 0; i < k; i++) result[i] = points[i];
        return result;
    }
}

Time: O(n log n). Space: O(1) nếu bạn tính sort là in-place (Timsort của Python dùng O(n) auxiliary; Array.Sort của C# dùng introsort với O(log n) stack depth). Dù sao đi nữa, không có cấu trúc dữ liệu phụ nào tỷ lệ với n.

Vấn đề ở đây: bạn sort những thứ bạn sẽ bỏ đi. Nếu n là 10,000 và k là 3, bạn đã thiết lập tổng thứ tự trên 10,000 item để lấy ra 3. Đó là công việc lãng phí, và đây chính là khoảng trống mà heap approach lấp đầy.

Approach 2: min heap (heapify rồi pop k lần)

Push tất cả điểm vào một min heap theo bình phương khoảng cách. Pop k lần. Mỗi lần pop lấy ra điểm có khoảng cách nhỏ nhất.

import heapq
 
class Solution:
    def kClosest(self, points: list[list[int]], k: int) -> list[list[int]]:
        heap = []
        for p in points:
            d = p[0] * p[0] + p[1] * p[1]
            heapq.heappush(heap, (d, p))
 
        result = []
        for _ in range(k):
            _, p = heapq.heappop(heap)
            result.append(p)
        return result

public class Solution {
    public int[][] KClosest(int[][] points, int k) {
        var minHeap = new PriorityQueue<int[], int>();
 
        foreach (var p in points) {
            int d = p[0] * p[0] + p[1] * p[1];
            minHeap.Enqueue(p, d);
        }
 
        int[][] result = new int[k][];
        for (int i = 0; i < k; i++) {
            result[i] = minHeap.Dequeue();
        }
        return result;
    }
}

Time: O(n log n) để xây heap, sau đó O(k log n) cho k lần pop. Tổng O(n log n). Space: O(n) — heap chứa tất cả n điểm.

Thời gian asymptotic giống với sorting. Bộ nhớ tệ hơn: O(n) heap so với O(1) sort. Không có lợi thế nào so với sorting cho bài toán này ở dạng đã nêu, trừ khi bạn đang stream các điểm từng cái một và không thể buffer tất cả chúng cho một sort — thì heap lại tự nhiên hơn. Với mảng tĩnh, sort thực sự tốt hơn ở đây.

Tôi chủ yếu dùng approach này khi giải thích cơ chế heap cho ai đó, vì cấu trúc rất rõ ràng: push mọi thứ vào, rồi pull ra k cái tốt nhất. Nó làm cho heap invariant trở nên dễ thấy.

Approach 3: bounded max heap (tối ưu khi k nhỏ so với n)

Cải tiến thực sự: thay vì một min heap kích thước n, giữ một max heap kích thước k. Invariant: heap luôn chứa k điểm gần nhất đã thấy cho đến nay, với điểm xa nhất trong số k điểm đó nằm ở đỉnh.

Khi gặp một điểm mới, nếu khoảng cách của nó nhỏ hơn giá trị lớn nhất của heap (điểm xa nhất trong k điểm hiện tại), loại bỏ điểm xa nhất và chèn điểm mới. Ngược lại bỏ qua điểm mới. Cuối cùng, heap chứa đúng k điểm gần nhất.

import heapq
 
class Solution:
    def kClosest(self, points: list[list[int]], k: int) -> list[list[int]]:
        # heapq của Python là min heap; đảo dấu distance để mô phỏng max heap
        max_heap = []
 
        for p in points:
            d = p[0] * p[0] + p[1] * p[1]
 
            if len(max_heap) < k:
                heapq.heappush(max_heap, (-d, p))
            elif d < -max_heap[0][0]:
                heapq.heapreplace(max_heap, (-d, p))
 
        return [p for _, p in max_heap]

public class Solution {
    public int[][] KClosest(int[][] points, int k) {
        // PriorityQueue của .NET là min heap; dùng comparer đảo ngược để làm max heap
        var maxHeap = new PriorityQueue<int[], int>(Comparer<int>.Create((a, b) => b.CompareTo(a)));
 
        foreach (var p in points) {
            int d = p[0] * p[0] + p[1] * p[1];
 
            if (maxHeap.Count < k) {
                maxHeap.Enqueue(p, d);
            } else if (d < maxHeap.TryPeek(out _, out int topPriority) ? topPriority : int.MaxValue) {
                maxHeap.Dequeue();
                maxHeap.Enqueue(p, d);
            }
        }
 
        return maxHeap.UnorderedItems.Select(x => x.Element).ToArray();
    }
}

Phiên bản C# với PriorityQueue cần một workaround vì PriorityQueue không expose priority hiện tại của phần tử đỉnh một cách sạch sẽ khi dùng như max heap. Một cách sạch hơn là tự quản lý heap:

public class Solution {
    public int[][] KClosest(int[][] points, int k) {
        var heap = new List<(int dist, int[] point)>(k);
        int heapCount = 0;
 
        foreach (var p in points) {
            int d = p[0] * p[0] + p[1] * p[1];
 
            if (heapCount < k) {
                heap.Add((d, p));
                heapCount++;
                if (heapCount == k) {
                    for (int i = k / 2 - 1; i >= 0; i--) SiftDown(heap, i, k);
                }
            } else if (d < heap[0].dist) {
                heap[0] = (d, p);
                SiftDown(heap, 0, k);
            }
        }
 
        return heap.Select(x => x.point).ToArray();
    }
 
    private void SiftDown(List<(int dist, int[] point)> heap, int i, int size) {
        while (true) {
            int largest = i;
            int left = 2 * i + 1, right = 2 * i + 2;
            if (left < size && heap[left].dist > heap[largest].dist) largest = left;
            if (right < size && heap[right].dist > heap[largest].dist) largest = right;
            if (largest == i) break;
            (heap[i], heap[largest]) = (heap[largest], heap[i]);
            i = largest;
        }
    }
}

Time: O(n log k). Mỗi trong n điểm kích hoạt tối đa một heap operation tốn O(log k). Space: O(k).

Trace tay qua ví dụ

points = [[1,3],[-2,2]], k = 1.

Heap giữ tối đa 1 phần tử. Khi gặp [-2,2] với khoảng cách 8, nó thắng so với maximum hiện tại (10), nên thay thế.

Với points = [[3,3],[5,-1],[-2,4]], k = 2:

Bình phương khoảng cách: [3,3] -> 18, [5,-1] -> 26, [-2,4] -> 20.

• Push [3,3]: heap = [(-18, [3,3])], kích thước 1 < 2.
• Push [5,-1]: heap = [(-26, [5,-1]), (-18, [3,3])], kích thước 2 = k. Root của max heap là -26, tức điểm xa nhất đang theo dõi là [5,-1] với khoảng cách 26.
• Gặp [-2,4]: khoảng cách 20. 20 < 26? Có. Thay root: loại [5,-1], push [-2,4]. Heap = [(-20, [-2,4]), (-18, [3,3])].
• Xong. Kết quả: [[3,3], [-2,4]].

Edge cases

k = 1. Cả ba approach đều hoạt động. Max heap giữ đúng một item và so sánh từng điểm tiếp theo với nó. Không có gì đặc biệt.

k = points.length. Max heap không bao giờ đủ đầy để evict bất kỳ thứ gì — mọi điểm đều được push. Cuối cùng chứa tất cả n điểm. Đúng.

points chứa [0,0]. Bình phương khoảng cách là 0. Nếu k >= 1, gốc tọa độ luôn có trong kết quả. Bounded heap sẽ chấp nhận nó ngay lập tức và nó không bao giờ bị evict (không điểm nào khác có khoảng cách nhỏ hơn). Cả ba approach xử lý điều này mà không cần special-case.

Hai điểm có cùng khoảng cách. Bài toán nói "đáp án được đảm bảo là duy nhất (ngoại trừ thứ tự)" — nghĩa là không có tie ở vị trí k nơi đổi một điểm cho điểm khác sẽ cho kết quả hợp lệ khác. Bạn không cần lo về distance-tie tại vị trí k.

So sánh

Tôi sẽ ship cái nào

Với mảng tĩnh ở constraint này (n <= 10^4), sort toàn bộ là lựa chọn đầu tiên của tôi. Chỉ hai dòng trong Python, dựa vào sort chuẩn đã được kiểm tra kỹ, không có trick đảo dấu nào phức tạp. Bất kỳ ai đọc code cũng hiểu ngay.

Nếu n là 10^7 và k là 50, tôi sẽ dùng bounded max heap. Sự khác biệt bộ nhớ — giữ 10 triệu điểm so với 50 — là thực sự. Và O(n log 50) so với O(n log 10^7) là một constant factor đáng kể trong thực tế.

Phiên bản min heap tôi dùng chủ yếu trong ngữ cảnh phỏng vấn để chứng minh rằng tôi hiểu cách xây dựng heap, hoặc khi input là một stream mà tôi không thể buffer tất cả trước.

Vị trí trong series

Các bài trước trong series này — Kth Largest Element và Top K Frequent Elements — đã thiết lập phản xạ heap: khi bài toán nói "k-th" hay "top k", hãy nghĩ đến priority queue. Bài này thêm chiều kích hình học (khoảng cách Euclidean làm khóa) và làm rõ trade-off giữa min heap và bounded max heap.

Pattern bounded max heap áp dụng trực tiếp cho:

• LeetCode 215 — Kth Largest Element in an Array: bounded max heap tương tự, nhưng khóa so sánh là giá trị chính nó chứ không phải khoảng cách tính toán. Solution O(n log k) ở đó là analogue trực tiếp.
• LeetCode 347 — Top K Frequent Elements: frequency thay thế khoảng cách; bạn xây max heap giới hạn kích thước k trên các cặp (frequency, element).
• LeetCode 692 — Top K Frequent Words: mở rộng 347 với tiebreak lexicographic — khóa so sánh trở thành (frequency, -lexicographic_order), đây là nơi trick đảo dấu xuất hiện theo hình thức ít rõ ràng hơn.
• LeetCode 378 — Kth Smallest Element in a Sorted Matrix: cấu trúc ma trận cho phép làm tốt hơn O(n log k), nhưng bounded heap là điểm khởi đầu hợp lý trước khi tiếp cận binary search.

Pattern chung ở tất cả: bạn có n item, muốn k item được xếp hạng theo một khóa nào đó, và bạn không cần tổng thứ tự đầy đủ. Hình dạng đó là bounded max heap.