K Closest Points to Origin — Day 8(Python)

Photo by park dasol on Unsplash

Today we will solve the k-closest neighbor problem. Before we look into the problem, we should know the formula for finding the distance between 2 points.

Distance formula. Source: Khan Academy

If you want to learn about the derivation of this formula, you can follow this link.

973. K Closest Points to Origin

We have a list of on the plane. Find the K closest points to the origin (0, 0).

(Here, the distance between two points on a plane is the Euclidean distance.)

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

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]]
(The answer [[-2,4],[3,3]] would also be accepted.)

Note:

1. 1 <= K <= points.length <= 10000
2. -10000 < points[i][0] < 10000
3. -10000 < points[i][1] < 10000

Solution 1:

One way to solve the problem is by finding the distance of each point from the origin and sorting the distance to find the answer. The code will look like the following.

class kClosestFinder:
    def kClosest(self, points: List[List[int]], K: int) -> List[List[int]]:
x2 = y2 = 0
        points = sorted(points, key = lambda x:math.sqrt(pow(x[0] - x2, 2) + pow(x[1] - y2, 2)))
        return points[:K]

I have tried generalizing the solution, which means you can change x2 or y2 if the problem requires you to do so. Now let us analyze the complexity of this solution.

Complexity analysis

Time Complexity

When we call the sorted function, timsort() is used behind the scene. The time complexity to perform timsort is O(NlogN).

Space Complexity

Timsort() uses O(N) space to perform the sort.

Can we try to improve the above logic?

How about we use the heap of size K and return the answer?

1. Run the loop for all the points in the list.
2. Find the distance of the current point from the origin.
3. Push it to the heap, if the size of the heap is equal to K, then pop from the heap before pushing.
4. Return the final heap.

import heapq
class kClosestFinder:
    def calculate_distance(self, x1, x2, y1, y2):
        return(math.sqrt(((x2-x1)**2) + (y1-y2)**2))
    
    def kClosest(self, points: List[List[int]], K: int) -> List[List[int]]:
        heap_list = []
        for (x, y) in points:
            dist = -self.calculate_distance(x, 0, y, 0)
            if len(heap_list) == K:
                heapq.heappushpop(heap_list, (dist, x, y))
            else:
                heapq.heappush(heap_list, (dist, x, y))
        
        return [(x,y) for (dist,x, y) in heap_list]

Time Complexity

To create the heap, we need to traverse the entire list that takes O(N). To build a heap of size K, the time complexity is O(log K).

Space Complexity

To build a heap of size K, it would take O(K).