K Closest Points to Origin

Here is a simple Python solution for this problem. We’ll use the built-in heapq module in Python, which allows us to create a max heap. The idea is to build a max heap of k elements and then for the remaining elements in the list, if they are closer than the farthest (i.e., max) element in the heap, we remove the max element and add the new one.

Here is the Python code:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

from typing import List
import heapq

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        # Create a max heap
        heap = []

        # Iterate over points
        for (x, y) in points:
            # Compute distance from origin
            dist = -(x*x + y*y)
            # If heap has less than k elements, we add the new point
            if len(heap) < k:
                heapq.heappush(heap, (dist, x, y))
            # Else if new point is closer than the farthest point in heap, we replace it
            elif dist > heap[0][0]:
                heapq.heappop(heap)
                heapq.heappush(heap, (dist, x, y))

        # Output only the points, excluding their distances
        return [(x,y) for (dist,x,y) in heap]

In this code, we start by creating an empty heap. Then we iterate over the given points. For each point, we calculate its distance from the origin (0,0), negating the value because Python’s heapq implements a min-heap and we want a max-heap. If our heap has less than k elements, we add the new point. Otherwise, we compare the distance of the current point with the largest distance in our heap. If the current point is closer, we remove the farthest point from the heap and add the current point. Finally, we return a list of the k closest points, discarding the distance information.