Heap

K Closest Points to Origin

medium

DESCRIPTION (inspired by Leetcode.com)

Given a list of points in the form [[x₁, y₁], [x₂, y₂], ... [x_n, y_n]] and an integer k, find the k closest points to the origin (0, 0) on the 2D plane.

The distance between two points (x, y) and (a, b) is calculated using the formula:

√(x₁ - a₂)² + (y₁ - b₂)²

Return the k closest points in any order.

Example 1:

Inputs:

points = [[3,4],[2,2],[1,1],[0,0],[5,5]]
k = 3

Output:

[[2,2],[1,1],[0,0]]

Also valid:

[[2,2],[0,0],[1,1]]
[[1,1],[0,0],[2,2]]
[[1,1],[2,2],[0,0]]
...
[[0,0],[1,1],[2,2]]

Code Editor

Run your code to see results here

Have suggestions or found something wrong?

Explanation

Approach 1: Sorting

The simplest approach is to sort calculate the distance of each point from the origin and sort the points based on their distance. This approach has a time complexity of O(n log n) where n is the number of points in the array, and a space complexity of O(n) (to store the sorted array of distances).

Approach 2: Max Heap

This problem can be solved using a similar approach to the one used to solve Kth Largest Element in an Array. The key difference is that we need to store the k closest points to the origin, rather than the k largest elements. Since we are looking for the k smallest elements, we need a max-heap, rather than a min-heap.

By default, python's heapq module implements a min-heap, but we can make it behave like a max-heap by negating the values of everything we push onto it.

We add the first k points to the heap by pushing a tuple containing the negative of the distance from the origin, and the index of the point. After that is finished, our heap contains the k closest points to the origin that we've seen so far, with the point furthest from the origin at the root of the heap.

k = 3

Visualization


def k_closest(points, k):
    heap = []
    for i in range(len(points)):
        x, y = points[i]
        distance = x * x + y * y
        
        if len(heap) < k:
            heapq.heappush(heap, (-distance, i))
        elif distance < -heap[0][0]:
            heapq.heappushpop(heap, (-distance, i))
    
    return [points[p[1]] for p in heap]

initialize heap

0 / 6

Pushing the first `k = 3` elements onto the heap. The root of the heap after this is done holds a tuple containing the negative of the distance of the point furthest from the origin, and the index of that point (-25, 0).

For each point after the first k, we calculate the distance from the origin and compare it with the root of the heap. If the current point is closer to the origin than the root of the heap, we pop the root and push the current point into the heap. This way, the heap will always contain the k closest points to the origin we've seen so far.

k = 3

Visualization


def k_closest(points, k):
    heap = []
    for i in range(len(points)):
        x, y = points[i]
        distance = x * x + y * y
        
        if len(heap) < k:
            heapq.heappush(heap, (-distance, i))
        elif distance < -heap[0][0]:
            heapq.heappushpop(heap, (-distance, i))
    
    return [points[p[1]] for p in heap]

push to heap

0 / 2

At the end of the iteration, the heap will contain the k closest points to the origin. We can iterate over each point in the heap and return the point associated with each tuple.

k = 3

Visualization


def k_closest(points, k):
    heap = []
    for i in range(len(points)):
        x, y = points[i]
        distance = x * x + y * y
        
        if len(heap) < k:
            heapq.heappush(heap, (-distance, i))
        elif distance < -heap[0][0]:
            heapq.heappushpop(heap, (-distance, i))
    
    return [points[p[1]] for p in heap]

distance = 50

0 / 1

Solution

Try these examples:

Visualization


def k_closest(points, k):
    heap = []
    for i in range(len(points)):
        x, y = points[i]
        distance = x * x + y * y
        
        if len(heap) < k:
            heapq.heappush(heap, (-distance, i))
        elif distance < -heap[0][0]:
            heapq.heappushpop(heap, (-distance, i))
    
    return [points[p[1]] for p in heap]

k closest points to origin

0 / 11

What is the time complexity of this solution?

O(1)

O(n log n)

O(n log k)

O(2ⁿ)

Mark as read

Next: Find K Closest Elements