[Leetcode] K Closest Points to Origin

heapq api in python makes this problem trivial.

Description

Given an array of points where points[i] = [xi, yi] represents a point on the X-Y plane and an integer k, return the k closest points to the origin (0, 0).

The distance between two points on the X-Y plane is the Euclidean distance (i.e., √(x1 - x2)2 + (y1 - y2)2).

You may return the answer in any order. The answer is guaranteed to be unique (except for the order that it is in).

Examples

Example 1:

Input: points = [[1,3],[-2,2]], k = 1
Output: [[-2,2]]
Explanation:
The distance between (1, 3) and the origin is sqrt(10).
The distance between (-2, 2) and the origin is sqrt(8).
Since sqrt(8) < sqrt(10), (-2, 2) is closer to the origin.
We only want the closest k = 1 points from the origin, so the answer is just [[-2,2]].

Example 2:

Input: points = [[3,3],[5,-1],[-2,4]], k = 2
Output: [[3,3],[-2,4]]
Explanation: The answer [[-2,4],[3,3]] would also be accepted.

Solution

Heap(priority queue) keeps the smallest(minheap) or the largest(maxheap) elem at top(root). The time complexity of heappush(push a new elem into heap and update the top) is O(h) or O(log n) To get the k closest, maintain the heap with no more than k elements. That is, pop largests elements when number of elems reaches k and push in new elem. By doing so, The heapify only takes O(log k). Traversing through the array is O(n). Hence, the total time complexity is O(n log k), faster than directly using sort of O(n log n)

To pop the largest elem when needed, the largest need to be at top, aka a maxheap is required. Note the heapq of python only provides minheap. Multiply the value by -1 so the point of max value owns min value and can be placed at top.

code

ref: https://medium.com/analytics-vidhya/k-closest-points-to-origin-day-8-python-4ee1e15f32d4