K Closest Points to Origin is a common question for Meta onsites. I've seen it a few times in the IQ discussion section recently, and the solution that they expect is a little unconventional, so here's my solution and explanation. You could solve it in O(NlgK) time using a min heap, but that's fairly straightforward, and it's probably not what most Meta interviewers are looking for.

I'm going to solve this using quickselect, which also seems to show up in Meta interviews a lot. O(N) average time, O(1) space, O(N^2) theoretical worst case time (but in practice we won't even get close to this). It took me a minute to understand why it's O(N) time on average, so I'll explain the solution and why it's O(N) on average below.

from random import randint

class Solution:
    def d(self, p):
	    return p[0]**2 + p[1]**2

    def swap(self, points, i, j):
        points[i], points[j] = points[j], points[i]
        
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        n = len(points)
        l, r = 0, n - 1
        while l != r:
            print(l, r)
            piv = randint(l, r)
            i, j = piv - 1, piv + 1
            while i >= l or j <= r:
			    if i >= 0 and self.d(points[i]) > self.d(points[piv]):
                    self.swap(points, i, piv-1)
                    self.swap(points, piv, piv-1)
                    piv -= 1
                i -= 1
                    
                if j < n and self.d(points[j]) < self.d(points[piv]):
                    self.swap(points, j, piv+1)
                    self.swap(points, piv, piv+1)
                    piv += 1
                j += 1
            
            if piv == k:
                break
            elif piv < k:
                l = piv
            else:
                r = piv - 1
                    
        return points[:k]

What's the logic behind the code? How does it solve the problem?

If you're familiar with this problem or you've read my solution, then you probably know that quickselect involves taking a single "pivot" element and iterating through the array to place all elements less than your pivot to the left, and place all elements greater than your pivot to the right.

Why is this useful for the K-closest points problem? For two reasons:

1. We can think of our K-closest points as sortable by their distance from the origin. In my code I use what I call "distance magnitude" which is just the distance formula without the square root. It saves a bit of time.
2. The K-closest points themselves don't need to be sorted. We just need to return the K points with the smallest distance magnitude, but they can be in any order.

The conclusion you should draw from the above two properties is that, if we select a point such that the pivot ends up being at index k, then we've got our answer because all indices 0 <= i < k will be the k points in the set with the lowest distance. So here's what we do:

1. Set our left and right bounds to be L = 0 and R = N - 1.
2. Select a pivot index in this range. You can technically select whatever pivot you want as long as it's in the interval [L, R], but I choose a random pivot. I'll explain why in the next section.
3. Iterate through the range [L, R] to place all elements less than our pivot to the left of it, and all elements greater than our pivot to the right of it.
4. If the index that our pivot ends up at is p, do the following:
   a. If p == k, break the iteration.
   b. If p < k, we want to get the pivot to the right, so set L = p.
   c. If p > k, we want to get the pivot to the left, so set R = p - 1.

Since the interval [L, R] gets smaller at each iteration and it always contains k, we will eventually find p == k here. Return points[:k] after this happens, because the first k points will definitely be the smallest after this happens.

Why is this solution O(N) time on average? The worst case is O(N^2), so is it actually better than a min-heap or sorting solution in practice?

Let's think about how we can get the worst case. At this point we have no reason to assume that the initial pivot index we choose matters (and for random input, it doesn't indeed), so let's say we always choose R as our initial pivot index. Then our worst case is if the input is sorted and k is 1; if so, [L, R] will look like this: [0, N - 1], [0, N - 2], ..., [0, 0]. We won't return until the very end after N iterations, and at each iteration, we look through roughly N elements. So the worst case time complexity is O(N^2), similar to if you had a nested loop from i = [0, N] and j = [i, N].

It's nice to assume that your lists' contents are random, but sorted lists are pretty common in practice, so let's protect ourselves from this worst case by selecting a random pivot index every time instead of the left- or right-most ones. If we do this, then the only way we could get the worst case is if we somehow manage to select the (N - i)th largest element as our pivot at each iteration i, so it becomes negligibly unlikely that we'll reach the worst case as N increases (because the probability of selecting a pivot that doesn't reduce [L, R] by much goes down as N goes up).

Indeed, if we choose a random pivot, then we'll end up roughly halving the size of [L, R] at each iteration. Recall that we iterate until L == R, and at each iteration, we look at (R - L) elements. So, if we approximately halve the array at each iteration, we should have about lg(N) iterations altogether. In the first iteration we'll check N elements, then N/2 elements, then N/4, etc. The sum of elements we're checking is sum(N/(2^i) for i in [0, lgN]), which is approximately 2*N. Therefore, in the average case, we take about O(2*N) ~= O(N) operations.

Notes

If this throws you for a bit of a loop, don't be discouraged. I think that the reason Meta asks questions like this is because it truly isn't obvious that this solution is better than one that has a faster higher TC bound, so it requires you to think outside of the box a little bit. The algorithm itself is somewhat challenging to write, too.

Feel free to ask any questions in the comments or point out if I made any mistakes.