K Closest Points to Origin
Bạn có n điểm trên mặt phẳng: điểm i là (xs[i], ys[i]). Trả về k điểm gần gốc tọa độ (0, 0) nhất, mỗi điểm dưới dạng một cặp [x, y].
Sắp xếp câu trả lời theo khoảng cách đến gốc tọa độ, điểm gần nhất trước. Nếu các điểm cách đều nhau, hãy xếp điểm có x nhỏ hơn trước, sau đó đến điểm có y nhỏ hơn. Một điểm xuất hiện nhiều lần được tính một lần cho mỗi bản sao.
Hàm
- xsinteger-array
- tọa độ x của mỗi điểm
- ysinteger-array
- tọa độ y của mỗi điểm
- kinteger
- cần trả về bao nhiêu điểm
- Trả vềinteger-2d-array
- các điểm gần nhất theo cặp [x, y], sắp xếp từ gần nhất trước, nếu bằng nhau thì theo x rồi đến y
Ràng buộc
n == xs.length == ys.length1 ≤ k ≤ n ≤ 5000-104 ≤ xs[i], ys[i] ≤ 104- Các điểm có thể lặp lại; mỗi bản sao là một điểm riêng.
Ví dụ
- Đầu vào
- xs = [3, -1, 2, 0]ys = [1, 2, -2, 5]k = 2
- Đầu ra
- [[-1, 2], [2, -2]]
- Giải thích
- Các khoảng cách bình phương lần lượt là 10, 5, 8 và 25. Hai khoảng cách nhỏ nhất thuộc về
(-1, 2)với giá trị 5 và(2, -2)với giá trị 8, theo thứ tự đó.
- Đầu vào
- xs = [1, -1, 0, 1]ys = [1, 1, 2, -1]k = 3
- Đầu ra
- [[-1, 1], [1, -1], [1, 1]]
- Giải thích
(1, 1),(-1, 1)và(1, -1)đều có khoảng cách bình phương bằng 2, còn(0, 2)bằng 4. Trong các điểm hòa nhau,(-1, 1)có giá trị x nhỏ nhất, còn(1, -1)đứng trước(1, 1)theo giá trị y.
- Đầu vào
- xs = [7]ys = [-7]k = 1
- Đầu ra
- [[7, -7]]
- Giải thích
- Chỉ có một điểm và bạn cần một điểm, nên đáp án chính là
(7, -7).
Câu hỏi mở rộng
Các điểm giờ đây lần lượt được đưa vào, và sau mỗi điểm, bạn phải báo cáo k điểm gần nhất tính đến lúc đó. Bạn cần bao nhiêu bộ nhớ?
Gợi ý
Mở từng gợi ý một. Mỗi gợi ý tiết lộ thêm một chút.
Bạn không bao giờ cần biết khoảng cách thực tế. Thay vì
√(x² + y²), bạn có thể so sánh đại lượng nào mà vẫn giữ nguyên thứ tự?Giả sử bạn giữ lại
kđiểm gần nhất đã thấy cho đến hiện tại. Một điểm mới sẽ cạnh tranh với điểm nào trong số đó?Duy trì một max-heap gồm tối đa
kđiểm, được sắp xếp theo(x² + y², x, y). Thay phần tử đầu heap khi có điểm mới đứng trước nó, rồi lấy các điểm ra khỏi heap ở cuối để điền đáp án từ phía sau.
Lời giải
Không cần tính căn bậc hai để tìm khoảng cách: x² + y² cho ra thứ tự như cũ và vẫn là số nguyên. Bạn chỉ cần k trong số n điểm, nên dùng một đống cực đại chứa k điểm gần nhất tính đến thời điểm hiện tại, với điểm xa nhất ở trên cùng để sẵn sàng thay thế, có độ phức tạp O(n log k) thay vì sắp xếp toàn bộ.
Chọn điểm gần nhất k lần
Đúng, nhưng không chạy xong với các test lớn nhất
Ý tưởng
Điểm đầu tiên trong đáp án là điểm gần nhất trong tất cả các điểm, điểm thứ hai là điểm gần nhất trong số những điểm còn lại, cứ tiếp tục như vậy. Vì thế, hãy quét các điểm k lần, mỗi lần chọn điểm đứng đầu theo thứ tự (x² + y², x, y) trong số những điểm chưa được chọn.
Mỗi lần quét đều xét cả n điểm. Với n = 5000 và k = 2500, ta có 1.25 × 10^7 phép so sánh ba thành phần.
Thuật toán
- Đánh dấu mọi điểm là chưa được chọn.
- Lặp lại
klần: quét các điểm chưa được chọn và tìm điểm đứng đầu theo thứ tự(x² + y², x, y). - Đánh dấu điểm đó là đã được chọn và thêm
[x, y]vào đáp án. - Trả về đáp án.
def kClosest(xs, ys, k):
n = len(xs)
taken = [False] * n
result = []
for _ in range(k):
best = -1
for i in range(n):
# Compare (squared distance, x, y): the order the answer uses.
if not taken[i] and (best == -1 or
(xs[i] * xs[i] + ys[i] * ys[i], xs[i], ys[i]) <
(xs[best] * xs[best] + ys[best] * ys[best], xs[best], ys[best])):
best = i
taken[best] = True
result.append([xs[best], ys[best]])
return resultSắp xếp mọi điểm
Ý tưởng
Đáp án là k điểm đầu tiên theo thứ tự đầy đủ, vì vậy hãy sắp xếp mọi điểm theo (x² + y², x, y) rồi giữ lại k điểm đầu tiên. Trong ví dụ 1, các khóa đã sắp xếp là (5, -1, 2), (8, 2, -2), (10, 3, 1) và (25, 0, 5), hai điểm đầu tiên tạo thành đáp án.
Việc sắp xếp có độ phức tạp O(n log n) bất kể k là bao nhiêu. Khi k nhỏ, phần lớn công việc đó dùng để sắp thứ tự các điểm không bao giờ nằm trong đáp án.
Thuật toán
- Tạo khóa cho từng điểm
(x² + y², x, y). - Sắp xếp các điểm theo khóa.
- Trả về
kđiểm đầu tiên dưới dạng các cặp[x, y].
def kClosest(xs, ys, k):
points = list(zip(xs, ys))
# Sort by squared distance, then x, then y: the order the answer uses.
points.sort(key=lambda p: (p[0] * p[0] + p[1] * p[1], p[0], p[1]))
return [[x, y] for x, y in points[:k]]Giữ k phần tử gần nhất trong một max-heap
Ý tưởng
Giữ lại k điểm tốt nhất đã gặp. Một điểm mới chỉ có thể loại điểm đang giữ có thứ hạng thấp nhất, vì vậy hãy lưu các điểm đang giữ trong một đống cực đại theo (x² + y², x, y), với điểm đó ở trên cùng. Khi đã lưu đủ k điểm, điểm nào có thứ hạng cao hơn điểm trên cùng sẽ thay thế điểm đó; các điểm còn lại bị loại, vì đã có k điểm được lưu có thứ hạng cao hơn chúng.
Trong ví dụ 2 với k = 3, điểm trên cùng sau ba điểm đầu tiên là (0, 2) với khóa (4, 0, 2). Điểm (1, -1) có khóa (2, 1, -1), có thứ hạng cao hơn và thay thế điểm đó.
Đống chỉ được sắp xếp một phần, vì vậy hãy lấy điểm trên cùng ra k lần và điền đáp án từ vị trí cuối cùng đến vị trí đầu tiên. Mỗi điểm tốn tối đa một bước đống O(log k), nên thời gian là O(n log k) và bộ nhớ là O(k).
Thuật toán
- Tạo một max-heap rỗng với khóa là
(x² + y², x, y). - Chèn từng điểm trong số
kđiểm đầu tiên. - Với mỗi điểm tiếp theo, so sánh điểm đó với phần tử trên cùng: nếu điểm đó xếp trước, hãy thay phần tử trên cùng bằng điểm đó rồi sàng xuống; nếu không thì bỏ qua.
- Lấy phần tử trên cùng ra
klần; mỗi điểm được lấy ra sẽ lấp vào vị trí trống cuối cùng trong đáp án. - Trả về đáp án.
import heapq
def kClosest(xs, ys, k):
# heapq is a min-heap; negating every part of (distance, x, y) puts the
# farthest kept point, in the answer's order, on top.
heap = []
for x, y in zip(xs, ys):
item = (-(x * x + y * y), -x, -y)
if len(heap) < k:
heapq.heappush(heap, item)
elif item > heap[0]:
# The new point comes before the farthest kept one: swap them.
heapq.heapreplace(heap, item)
# Popping gives the farthest first, so fill the answer from the back.
result = [None] * len(heap)
for i in range(len(result) - 1, -1, -1):
_, x, y = heapq.heappop(heap)
result[i] = [-x, -y]
return result
Lỗi thường gặp và trường hợp biên
- Bỏ qua bước phân định thứ tự khi hòa. Ví dụ 2 có ba điểm cùng cách gốc tọa độ 2 đơn vị; nếu không so sánh
xrồi đếny, chúng sẽ giữ nguyên thứ tự đầu vào thay vì theo thứ tự(-1, 1),(1, -1),(1, 1). - Dùng heap min chứa
kđiểm. Phần tử trên cùng là điểm gần nhất, nên mỗi lần loại bỏ sẽ bỏ đi một điểm cần có trong đáp án, và cuối cùng bạn sẽ cókđiểm xa nhất. - Đọc thẳng đáp án từ mảng của heap. Mảng đó chỉ được sắp xếp một phần; trước tiên, hãy lần lượt lấy từng điểm ra.
- Loại bỏ các điểm trùng lặp. Nếu điểm
(5, 5)xuất hiện ba lần vàk = 3, đáp án là điểm đó ba lần.
Câu hỏi thường gặp4
Độ phức tạp thời gian của bài toán tìm K điểm gần gốc tọa độ nhất là bao nhiêu?
Sắp xếp tất cả các điểm mất O(n log n). Phiên bản dùng heap có độ phức tạp O(n log k) cho bước quét, cộng thêm O(k log k) để sắp xếp k điểm còn lại, và chỉ giữ O(k) điểm trong bộ nhớ.
Tại sao dùng đống cực đại thay vì đống cực tiểu để tìm k điểm gần nhất?
Một điểm mới cạnh tranh với điểm xa nhất trong số k điểm được giữ lại, vì vậy điểm đó phải nằm trên cùng. Min-heap chỉ hoạt động nếu nó chứa tất cả n điểm để bạn có thể xóa k điểm gần nhất, việc này tốn O(n) bộ nhớ thay vì O(k).
Bạn có cần căn bậc hai để so sánh khoảng cách không?
Không. Với các số không âm, a < b khi và chỉ khi √a < √b, nên x² + y² cho thứ tự giống như khoảng cách. Giá trị này cũng luôn là số nguyên: tối đa là 2 × 10^8 trong trường hợp này, nằm trong phạm vi 32-bit.
Có thể giải bài toán K điểm gần gốc tọa độ nhất trong thời gian O(n) không?
Trung bình là có. Quickselect phân vùng các điểm quanh một trục chọn ngẫu nhiên theo cùng thứ tự (x² + y², x, y) và chỉ tiếp tục ở phía chứa vị trí k, với độ phức tạp trung bình là O(n). k điểm nằm bên trái vị trí đó là đáp án; việc sắp xếp chúng tốn thêm O(k log k).
Bài tập tương tự
Các bài tập dùng cùng ý tưởng. Giải hai hoặc ba bài là bạn sẽ nhớ được mẫu này.
def kClosest(xs, ys, k):
# Viết code ở đâyKiểm thử #1
Kiểm thử #2
Kiểm thử #3
Đối số
xs = [3, -1, 2, 0] ys = [1, 2, -2, 5] k = 2
Đầu ra mong đợi
[[-1, 2], [2, -2]]