Add 'K Closest Points to Origin'

Author: nickolashkraus

README.md

@@ -20,6 +20,8 @@ A collection of LeetCode solutions
 
 [Invert Binary Tree](./src/invert_binary_tree.py)
 
+[K Closest Points to Origin](./src/k_closest_points_to_origin.py)
+
 [Linked List Cycle](./src/linked_list_cycle.py)
 
 [Longest Common Subsequence](./src/longest_common_subsequence.py)

src/k_closest_points_to_origin.py

@@ -0,0 +1,34 @@
+"""
+973. K Closest Points to Origin
+
+https://leetcode.com/problems/k-closest-points-to-origin
+
+NOTES
+  * I call it house work, 'cause it's light work. Calculate distance,
+    maintain max-heap (closest distance is equivalent to smallest distance) of
+    k points.
+"""
+
+import heapq
+import math
+
+
+class Solution:
+    """
+    Maintains a max-heap of size k containing the k closest points to the
+    origin. Since the `heapq` module only provides an implementation for a min-
+    heap, the distance is negated before being inserted into the heap.
+    """
+    def kClosest(self, points: list[list[int]], k: int) -> list[list[int]]:
+
+        heap: list[tuple[float, list[int]]] = []  # (distance, [xi,yi])
+
+        for i, p in enumerate(points):
+            d = math.sqrt(p[0]**2 + p[1]**2)
+            if i < k:
+                heapq.heappush(heap, (-d, p))
+            else:
+                if -d > heap[0][0]:
+                    heapq.heapreplace(heap, (-d, p))
+
+        return [p[1] for p in heap]

tests/test_k_closest_points_to_origin.py

@@ -0,0 +1,19 @@
+"""
+973. K Closest Points to Origin
+
+https://leetcode.com/problems/k-closest-points-to-origin
+"""
+
+from unittest import TestCase
+
+from src.k_closest_points_to_origin import Solution
+
+
+class TestSolution(TestCase):
+    def test_1(self):
+        exp = sorted([[-2, 2]])
+        assert sorted(Solution().kClosest([[1, 3], [-2, 2]], 1)) == exp
+
+    def test_2(self):
+        exp = sorted([[3, 3], [-2, 4]])
+        assert sorted(Solution().kClosest([[3, 3], [5, -1], [-2, 4]], 2)) == exp