973. K Closest Points

Problem Statement in English

You’re given a list of cordinates, points. Your task is to return the $k$ points closest to the origin.

Approach

We just need to maintain a min heap, and return the first $k$ values in it - essentially the $k$ smallest values.

Solution in Python

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        minHeap = []

        for cordinate in points:
            dist = cordinate[1] ** 2 + cordinate[0] ** 2
            minHeap.append([dist, *cordinate])

        heapq.heapify(minHeap)

        res = []

        while k > 0:
            _, x, y = heapq.heappop(minHeap)
            res.append([x, y])
            k -= 1

        return res

*cordinate is just list unpacking in Python.

Note 1: The reason we place dist at the beginning of each list is so that Python uses that to sort the elements.

To further optimise the number of elements in the min heap, and improve runtime, we could have made it a max heap and then maintained only the $k$ smallest values. That would look like this.

class Solution:

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

        for cordinate in points:
            dist = cordinate[1] ** 2 + cordinate[0] ** 2
            heapq.heappush(minHeap, [-dist, *cordinate])

            if len(minHeap) > k:
                heapq.heappop(minHeap)

        res = []

        while k > 0:
            _, x, y = heapq.heappop(minHeap)
            res.append([x, y])
            k -= 1

        return res

The reason I multiplied dist by $-1$ is because there’s no other way to create a max heap using the built-in heapq module.

We will consider this approach in the complexity calculation.

Complexity

• Time: $O(nlog(k))$
  Since we only maintain $k$ elements in the max heap. And to insert $n$ elements in a heap of size $k$, the time complexity becomes $O(nlog(k))$.

• Space: $O(k)$
  Since we only store the $k$ smallest values.

Mistakes I Made

It took a while for me to figure out that heapq doesn’t give you a sorted list, and that you need to use the API it provides to properly work with the heap.

Figuring out how to use the max heap was very traumatising. I’m kind of disappointed there isn’t a more elegant option in Python.

And we are done.