K Closest Points to Origin

Data Structures and Algorithms

Statement

Given a list of points on a plane, where the plane is a 2-D array with (x, y) coordinates, find the k closest points to the origin (0, 0).

Constraints

• 1 ≤ k ≤ points.length() ≤ 10³
• -10⁴ ≤ x[i], y[i] ≤ 10⁴

Solution

"""
production algorithm
"""

from math import sqrt
from data_structures.heap.heap import Heap


class Solution:
    def k_closest(self, points, k):
        """
        time complexity O(nlogk)
        space complexity O(k)
        """
        closest = Heap()
        for i in range(len(points)):
            closest.push(MaxHeapData(self.distance(points[i]), points[i]))
            if k < closest.size():
                closest.pop()

        result = list()
        for i in range(k):
            _, point = closest.pop()
            result.append((point.x, point.y))

        return result

    def distance(self, point):
        """
        time complexity O(1)
        space complexity O(1)
        """
        return sqrt(point.x ** 2 + point.y ** 2)


class Point:
    def __init__(self, x, y):
        self.x = x
        self.y = y


class MaxHeapData:
    def __init__(self, distance, point):
        self.distance = distance
        self.point = point

    def __lt__(self, other):
        return other.distance < self.distance

    def __iter__(self):
        yield self.distance
        yield self.point

    def __len__(self):
        return 2

"""
unit tests
"""

from unittest import TestCase
from algorithms.top_k_elements.k_closest_points_to_origin.solution import Solution, Point


class SolutionTestCase(TestCase):
    def test_points_on_unit_square(self):
        # Given
        points = [Point(1, 0), Point(1, 1), Point(0, 1), Point(-1, 1),
                  Point(-1, 0), Point(-1, -1), Point(0, -1), Point(1, -1)]
        k = 4
        solution = Solution()

        # When
        closest = solution.k_closest(points, k)
        actual = set(closest)

        # Then
        expected = set([(1, 0), (0, 1), (-1, 0), (0, -1)])
        self.assertEqual(actual, expected)

    def test_points_on_unit_square_plus_origin(self):
        # Given
        points = [Point(0, 0), Point(1, 0), Point(1, 1), Point(0, 1), Point(-1, 1),
                  Point(-1, 0), Point(-1, -1), Point(0, -1), Point(1, -1)]
        k = 4
        solution = Solution()

        # When
        closest = solution.k_closest(points, k)
        actual = set(closest)

        # Then
        expected = set([(0, 0), (1, 0), (0, 1), (0, -1)])
        self.assertEqual(actual, expected)

Strategy

Top K Elements.

Explanation

Sort all points in a max heap. When the max heap reaches size k+1, then pop the top of the max heap. At the top of the max heap is the kth closest point to the origin.

Time Complexity

The max heap has size proportional to k. There are n iterations of operations on the max heap. Therefore, the time complexity of the algorithm is O(nlogk).

Space Complexity

The maximum size of the max heap is k. Therefore, the auxiliary space complexity of the algorithm is O(k).

Links

data_structures_algorithms/algorithms/top_k_elements/k_closest_points_to_origin at main ·…