Given a list of points in the form [[x1, y1], [x2, y2], ... [xn, yn]]
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:
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: Min Heap
This problem can be solved using a similar approach to the one used to solve Kth Largest Element in an Array. They 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
defk_closest(points, k):
heap =[]
for i inrange(len(points)):
x, y = points[i]
distance = x * x + y * y
iflen(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
Python
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
defk_closest(points, k):
heap =[]
for i inrange(len(points)):
x, y = points[i]
distance = x * x + y * y
iflen(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
Python
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
defk_closest(points, k):
heap =[]
for i inrange(len(points)):
x, y = points[i]
distance = x * x + y * y
iflen(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
Python
Solution
defk_closest(points, k):
heap =[]
for i inrange(len(points)):
x, y = points[i]
distance = x * x + y * y
iflen(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
Python
Test Your Knowledge
Login to take the complexity quiz and track your progress
Complexity Analysis
Time Complexity: O(n log k) where n is the number of points in the array and k is the input parameter. We iterate over all points, and in the worst case, we both push and pop each point from the heap, which takes O(log k) time per point.
Space Complexity: O(k) where k is the input parameter. The space is used by the heap to store the k closest points to the origin.
Your account is free and you can post anonymously if you choose.