K Closest Points to Origin

Patrick Leaton
Published: June 10, 2021, 9:31 p.m.
Problem 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).

 

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 <= 10^4
  • -10^4 < xi, yi < 10^4

 

The description was taken from https://leetcode.com/problems/k-closest-points-to-origin/.

Problem Solution

#O(N(Log N)) Time, O(N) Space
class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        heap, output = [], []
       
        for x_one, y_one in points:
            heapq.heappush(heap, (math.sqrt((x_one)**2 + (y_one)**2), ([x_one,y_one])))
       
        while k > 0:
            distance, coordinates = heapq.heappop(heap)
            output.append(coordinates)
            k -= 1
       
        return output

Problem Explanation

Right off the bat, the problem title states we are returning the k least distant points from the center of the graph.

The "k" is usually an indication that this can be solved by utilizing a heap because it implies that we will be working off some sorted order.  The "least" tells us that this may be a min heap.

From those clues, the implementation is how we may expect.  We will use the Euclidean Distance formula that is given to us from the description to calculate the distance from the center point to each other point the input then push those distances into a heap.  We will then continuously pop the least distant point off of the heap and append it to the output.  Once we have done that k times, we are done.

Let's start by initializing our output and heap.

        heap, output = [], []

 

Next, we will calculate the difference between the two given points.

We are given this formula: √(x1 - x2)^2 + (y1 - y2)^2.  This formula just calculates the distance between two points on a one-dimensional plane by inputting both points' coordinates.  Basically, this is how we can draw a single straight line between two points.

Since the second point we are calculating the distance to is [0,0], we can simplify the statement to this: √(x1)^2 + (y1)^2. 

Let's traverse the array of points, and for each point coordinates we encounter within each iteration, we will calculate the distance based on the formula mentioned previously and push that, as well as the coordinates themselves into our heap.

        for x_one, y_one in points:
            heapq.heappush(heap, (math.sqrt((x_one)**2 + (y_one)**2), ([x_one,y_one])))

 

It is important the distance is first because the heap sorts on the first value inside the tuple.

Then, while we haven't popped k values off of our heap, we will continue to include the least current point distance.

        while k > 0:
            distance, coordinates = heapq.heappop(heap)
            output.append(coordinates)
            k -= 1

 

Once we have broken from the loop, we have found k least distances from the origin, so we will return them in the output.

        return output