Problem Statement

Description

Given an array of points where points[i] = [xi, yi] represents a point on the X-Y plane and an integer k, return the k closest points to the origin (0, 0).

The distance between two points on the X-Y plane is the Euclidean distance (i.e., √(x1 - x2)2 + (y1 - y2)2).

You may return the answer in any order. The answer is guaranteed to be unique (except for the order that it is in).

Examples

Example 1:

Input: points = [[1,3],[-2,2]], k = 1
Output: [[-2,2]]
Explanation:
The distance between (1, 3) and the origin is sqrt(10).
The distance between (-2, 2) and the origin is sqrt(8).
Since sqrt(8) < sqrt(10), (-2, 2) is closer to the origin.
We only want the closest k = 1 points from the origin, 
so the answer is just [[-2,2]].

Example 2:

Input: points = [[3,3],[5,-1],[-2,4]], k = 2
Output: [[3,3],[-2,4]]
Explanation: The answer [[-2,4],[3,3]] would also be accepted.

Constraints

  • 1 <= k <= points.length <= 104
  • -104 < xi, yi < 104

Key Insights

This is a pretty straight forward problem. It’s listed as a medium on Leetcode, but I think this one is probably more of an easy if you know how to perform sorting in Python or if you know how to use the heapq library.

The easiest approach is to use *.sort(...) to order the list by distance with a lambda function. Alternatively, the coordinates and their associated distances could be added to a min heap as tuples and the k closest coordinates could be popped from the heap and returned.

Python Solution

Sorting Solution

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        """
            T: O(N log N )
            S: O(1)

            - The euclidean distance requires a square root. 
              However, since we only care about the sort order,
              we can save an operation by not taking the square
              root of the result.
            - Notice how we can use the key = lambda x : ... 
              to create a custom sort function
            - Once we have sorted the list, we can return the
              k - closest points to the origin
              
        """
        points.sort(key = lambda x: x[0]**2 + x[1]**2)
        return points[:k]

Heap Solution

from heapq import heappush, heappop, heapify

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        """
            T: O( log N)
            S: O( N )

            - This is not the most efficient solution, but an interviewer
              could ask you use a heap. Using *.sort(...) will execute faster
              due to the underlying implementation
            - Loop over the points list, calculate the distance, and 
              push a distance, coordinate tuple onto the min heap
            - Pop the first k elements off of the heap and return them
              as a list.
        """
        distances = []
        heapify(distances)

        for x, y in points:
            dist = (x - 0)**2 + (y - 0)**2
            heappush(distances,  (dist, [x, y]))
        
        out = []

        for i in range(k):
            _, point = heappop(distances)
            out.append(point)

        return out