K Closest Points to Origin
DESCRIPTION (credit Leetcode.com)
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:
√(x1 - a2)2 + (y1 - b2)2
Return the k closest points in any order.
EXAMPLES
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]]
Run your code to see results here
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: 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.
import heapqdef k_closest(points, k):heap = []for i in range(len(points)):x, y = points[i]distance = x * x + y * yif 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]
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.
import heapqdef k_closest(points, k):heap = []for i in range(len(points)):x, y = points[i]distance = x * x + y * yif 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]
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.
import heapqdef k_closest(points, k):heap = []for i in range(len(points)):x, y = points[i]distance = x * x + y * yif 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]
Solution
import heapqdef k_closest(points, k):heap = []for i in range(len(points)):x, y = points[i]distance = x * x + y * yif 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]
Complexity Analysis
Time Complexity: O(n log k). We iterate over all the points in the array. Comparing the current point with the root of the heap takes O(1) time. In the worst case, we both push and pop each point from the heap, which takes O(log k) time.
Space Complexity: O(k). The space used by the heap to store the k closest points to the origin.
Loading comments...