K Closest Points to Origin

Max-heap

Problem

We have a list of points on the plane. Find the K closest points to the origin (0, 0).

(Here, the distance between two points on a plane is the Euclidean distance.)

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

For example:

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]].

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

Thought Process

• Since we are looking for the K closest points to the origin, we need to know what the furthest points from the origin are so we can discard them in place for the closest points (because our for-loop will skip any points/their distance that may be greater i.e. further away from the origin than the other K elements.) This is why we use a max heap.

• We use the Euclidean distance formula to calculate the distance.

• To associate points with their distances, in the max heap we will store the distance with their points.

• In each iteration of our for-loop from K we are saying: "Ok, here is the furthest point from the origin we have thus far in our heap, is the current distance smaller than this? If so, then pop this furthest point at heap[0] and insert this new point."

Solution

from heapq import *
class Solution:
    def kClosest(self, points: List[List[int]], K: int) -> List[List[int]]:
        heap = []
        
        for i in range(K):
            dist = self.eDist(points[i][0], points[i][1])
            heappush(heap,(-dist, points[i]))
            
        for i in range(K, len(points)):
            pointDist = self.eDist(points[i][0],points[i][1])
            if -(pointDist) > heap[0][0]:
                heappop(heap)
                heappush(heap,(-pointDist, points[i]))
        res = []
        for i in range(len(heap)):
            res.append(heap[i][1])
        return res
    
    def eDist(self, x, y):
        return ((x*x) + (y*y))

Key Points

• Using a max heap in order to know what the furthest points from the origin (large distances) are so we can discard them

• Storing the distances and points together in pairs in the heap

Time Complexity

• Time: $O(N*log(K))$ as we are iterating points and pushing them into the heap

• Space: $O(K)$