K Closest Points to Origin

The problem description can be found in LeetCode.

Approach

The easiest solution would be to sort the input and than take the first k elements. This would require $O(nlogn)$ time for the sorting and $O(k)$ for selecting values.

However, we can do better! In fact, we can use a Heap to store all the elements and then perform k pop operations to get the closest values to origin.

Solution

import heapq
from typing import List


class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        heap = [(x**2 + y**2, x, y) for x, y in points]
        heapq.heapify(heap)

        ans = []
        for _ in range(k):
            _, x, y = heapq.heappop(heap)
            ans.append([x, y])

        return ans

Complexity

• Time: $O(n+klogn)$
• Space: $O(n)$

Where $n$ is the size of the points list.