Add k closest points to origin (keon#435)

* Add k closest points to origin

* Update heap init file

* Add tests for k closest points to origin

* Add k closest points link to README

* Update k closest points to origin

Co-Authored-By: vinceajcs <[email protected]>

README.md

@@ -158,6 +158,7 @@ If you want to uninstall algorithms, it is as simple as:
     - [skyline](algorithms/heap/skyline.py)
     - [sliding_window_max](algorithms/heap/sliding_window_max.py)
     - [binary_heap](algorithms/heap/binary_heap.py)
+    - [k_closest_points](algorithms/heap/k_closest_points.py)
 - [linkedlist](algorithms/linkedlist)
     - [add_two_numbers](algorithms/linkedlist/add_two_numbers.py)
     - [copy_random_pointer](algorithms/linkedlist/copy_random_pointer.py)

algorithms/heap/__init__.py

@@ -1,3 +1,5 @@
 from .binary_heap import *
 from .skyline import *
 from .sliding_window_max import *
+from .merge_sorted_k_lists import *
+from .k_closest_points import *

algorithms/heap/k_closest_points.py

@@ -0,0 +1,44 @@
+"""Given a list of points, find the k closest to the origin.
+
+Idea: Maintain a max heap of k elements.
+We can iterate through all points.
+If a point p has a smaller distance to the origin than the top element of a heap, we add point p to the heap and remove the top element.
+After iterating through all points, our heap contains the k closest points to the origin.
+"""
+
+
+from heapq import heapify, heappushpop
+
+
+def k_closest(points, k, origin=(0, 0)):
+    # Time: O(k+(n-k)logk)
+    # Space: O(k)
+    """Initialize max heap with first k points.
+    Python does not support a max heap; thus we can use the default min heap where the keys (distance) are negated.
+    """
+    heap = [(-distance(p, origin), p) for p in points[:k]]
+    heapify(heap)
+
+    """
+    For every point p in points[k:],
+    check if p is smaller than the root of the max heap;
+    if it is, add p to heap and remove root. Reheapify.
+    """
+    for p in points[k:]:
+        d = distance(p, origin)
+
+        heappushpop(heap, (-d, p))  # heappushpop does conditional check
+        """Same as:
+            if d < -heap[0][0]:
+                heappush(heap, (-d,p))
+                heappop(heap)
+
+        Note: heappushpop is more efficient than separate push and pop calls.
+        Each heappushpop call takes O(logk) time.
+        """
+
+    return [p for nd, p in heap]  # return points in heap
+
+
+def distance(point, origin=(0, 0)):
+    return (point[0] - origin[0])**2 + (point[1] - origin[1])**2

tests/test_heap.py

@@ -1,15 +1,18 @@
 from algorithms.heap import (
     BinaryHeap,
     get_skyline,
-    max_sliding_window
+    max_sliding_window,
+    k_closest_points
 )
 
 import unittest
 
+
 class TestBinaryHeap(unittest.TestCase):
     """
         Test suite for the binary_heap data structures
     """
+
     def setUp(self):
         self.min_heap = BinaryHeap()
         self.min_heap.insert(4)
@@ -24,31 +27,41 @@ def test_insert(self):
         # After insert:  [0, 2, 50, 4, 55, 90, 87, 7]
         self.min_heap.insert(2)
         self.assertEqual([0, 2, 50, 4, 55, 90, 87, 7],
-                          self.min_heap.heap)
+                         self.min_heap.heap)
         self.assertEqual(7, self.min_heap.currentSize)
 
     def test_remove_min(self):
         ret = self.min_heap.remove_min()
         # Before remove_min : [0, 4, 50, 7, 55, 90, 87]
         # After remove_min: [7, 50, 87, 55, 90]
         # Test return value
-        self.assertEqual(4,ret)
+        self.assertEqual(4, ret)
         self.assertEqual([0, 7, 50, 87, 55, 90],
-                          self.min_heap.heap)
+                         self.min_heap.heap)
         self.assertEqual(5, self.min_heap.currentSize)
 
+
 class TestSuite(unittest.TestCase):
     def test_get_skyline(self):
-        buildings = [ [2, 9, 10], [3, 7, 15], [5, 12, 12], \
-                      [15, 20, 10], [19, 24, 8] ]
+        buildings = [[2, 9, 10], [3, 7, 15], [5, 12, 12],
+                     [15, 20, 10], [19, 24, 8]]
         # Expect output
-        output = [ [2, 10], [3, 15], [7, 12], [12, 0], [15, 10], \
-                   [20, 8], [24, 0] ]
+        output = [[2, 10], [3, 15], [7, 12], [12, 0], [15, 10],
+                  [20, 8], [24, 0]]
         self.assertEqual(output, get_skyline(buildings))
 
     def test_max_sliding_window(self):
         nums = [1, 3, -1, -3, 5, 3, 6, 7]
         self.assertEqual([3, 3, 5, 5, 6, 7], max_sliding_window(nums, 3))
 
+    def test_k_closest_points(self):
+        points = [(1, 0), (2, 3), (5, 2), (1, 1), (2, 8), (10, 2), (-1, 0), (-2, -2)]
+        self.assertEqual([(-1, 0), (1, 0)], k_closest(points, 2))
+        self.assertEqual([(1, 1), (-1, 0), (1, 0)], k_closest(points, 3))
+        self.assertEqual([(-2, -2), (1, 1), (1, 0), (-1, 0)], k_closest(points, 4))
+        self.assertEqual([(10, 2), (2, 8), (5, 2), (-2, -2), (2, 3),
+                          (1, 0), (-1, 0), (1, 1)], k_closest(points, 8))
+
+
 if __name__ == "__main__":
     unittest.main()