k closest point to origin

I am trying to solve this problem. https://leetcode.com/problems/k-closest-points-to-origin/ I have used maxheap and following is my code. This works well for points= [[1,3],[-2,2]], k = 1. However for points = [[3,3],[5,-1],[-2,4]], k = 2, this is giving the output as [[3,3],[5,-1]] whereas the expected output is [[3,3],[-2,4]]. I am not being able to figure out why my code is picking the point [5,-1]. Please help find the mistake I am doing.

from heapq import *
class Solution:

def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:

    maxHeap = []
    for i in range(k):
        heappush(maxHeap, points[i])
        
    for i in range(k, len(points)):
        dist = self.distance(points[i])
        heapdistance_from_origin = self.distance(maxHeap[0])
        if dist < heapdistance_from_origin:
            heappop(maxHeap)
            heappush(maxHeap, points[i])
    return maxHeap

def distance(self, point: List[int]):
     return point[0] * point[0] + point[1] * point[1]

Solution

The problem: heappush and heappop don't order the heap elements by distance to origin. They use the natural order of the elements instead (they're lists, so it's lexicographical order). So maxHeap[0] is usually not the point furthest from the origin.

You can use heapq.nsmallest, that supports a key argument (and it uses almost the same algorithm you do):

    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        return nsmallest(k, points, key=lambda p: hypot(*p))

Or:

    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        return nsmallest(k, points, key=partial(dist, (0, 0)))

Btw no need to import it yourself, LeetCode imports pretty much everything for you already. (Of course you can still do it if you want for clarity or so, then in my case it would be clear that hypot and dist come from math and that partial comes from functools. LeetCode is supposed to be about coding interviews, and I guess in a whiteboard interview, it's convenient to just say out loud "partial from functools" and have less to write, maybe that's why LeetCode does this.)

If you want to implement it yourself: Instead of making your heap hold bare points, you could make it hold pairs of (negated_distance, point). Negated because heapq implements a min-heap, so if you want a max-heap, negate the values. Here's that:

    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:

        maxHeap = []
        for p in points[:k]:
            heappush(maxHeap, (-self.distance(p), p))

        for p in points[k:]:
            dist = self.distance(p)
            heapdistance_from_origin = -maxHeap[0][0]
            if dist < heapdistance_from_origin:
                heappop(maxHeap)
                heappush(maxHeap, (-dist, p))
        return [p for _, p in maxHeap]

Or you could convert the points to instances of a Point class that does the comparison you want:

    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:

        class Point(list):
            def __lt__(self_, other):
                return self.distance(self_) > self.distance(other)
        points = list(map(Point, points))
        
        maxHeap = []
        for i in range(k):
            heappush(maxHeap, points[i])

        for i in range(k, len(points)):
            dist = self.distance(points[i])
            heapdistance_from_origin = self.distance(maxHeap[0])
            if dist < heapdistance_from_origin:
                heappop(maxHeap)
                heappush(maxHeap, points[i])
        return maxHeap