https://leetcode.com/problems/k-closest-points-to-origin

When you see in the statement of the problem that you need you find k biggest or k smallest elements, you should immediately think about heaps or sort. Here we need to find the k smallest elements, and hence we need to keep max heap. Why max and not min? We always keep in the root of our heap the k-th smallest element. Let us go through example: points = [[1,2],[2,3],[0,1]], [3,4], k = 2.

  1. First we put points [1,2] and [2,3] into our heap. In the root of the heap we have maximum element [2,3]
  2. Now, we see new element [0,1], what should we do? We compare it with the root, see, that it is smaller than root, so we need to remove it from our heap and put new element instead, now we have elements [1,2] and [0,1] in our heap, root is [1,2]
  3. Next element is [3,4] and it is greater than our root, it means we do not need to do anything.

Complexity

We process elements one by one, there are n of them and push it into heap and pop root from our heap, both have O(log k) complexity, so finally we have O(n log k) complexity, which is faster than O(n log n) algorighm using sorting.

Code:

First, we create heap (which is by definition min heap in python, so we use negative distances). Then we visit all the rest elements one by one and update our heap if needed.

class Solution:
    def kClosest(self, points, K):
        heap = [[-i*i-j*j, i, j] for i,j in points[:K]]
        rest_elem = [[-i*i-j*j, i, j] for i,j in points[K:]]
        heapq.heapify(heap)
        for s, x, y in rest_elem:
            if s > heap[0][0]:
                heapq.heappush(heap, [s,x,y])
                heapq.heappop(heap)

        return [[x,y] for s,x,y in heap]

Solution 2:

Even though this algorighm has not optimal algorithmic complexity (it is O(n log n) vs heaps O(n log k), on leetcode it can work faster. Just sort points by distances and choose the smallest K of them

class Solution:
    def kClosest(self, points, K):
        return sorted(points, key = lambda x: x[0]**2 + x[1]**2)[:K]

If you like the solution, you can upvote it on leetcode discussion section: Problem 0973