K Closest Points to Origin

Table of Contents +

You have many points on a map. You want the k points nearest to the center. Sorting them all works but wastes effort. You only need k of them. A heap of size k gives you those k and ignores the rest. The interviewer wants to see if you can avoid doing more work than needed.

🎯 The Problem

You get a list of points and a number k, and you return the k points closest to the origin (0, 0).

What you get:

  • A list of points. Each point has an x and a y value.
  • A number k for how many closest points to return.
  • Order among the k does not matter.

The distance rule:

  • The Euclidean distance to the origin is the square root of x*x + y*y.
  • But we do not need the square root. Comparing x*x + y*y gives the same order.
  • So we skip the square root and save work.
Input:  points = [[1, 3], [-2, 2], [5, 8], [0, 1]], k = 2
Output: [[0, 1], [-2, 2]]


Explanation:
distance squared: [1,3]=10, [-2,2]=8, [5,8]=89, [0,1]=1
The two smallest are 1 and 8 -> points [0,1] and [-2,2]

So we want the k points with the smallest distance. Order among the k does not matter.

Here is the idea drawn as points around the origin. We pick the closest k.

Origin (0,0)

(0,1) dist2 = 1

(-2,2) dist2 = 8

(1,3) dist2 = 10

(5,8) dist2 = 89

🐢 Approach 1: Sort All Points (Brute Force)

The idea in one line: measure every point, sort by distance, take the first k.

The idea:

  • Compute the distance squared of every point.
  • Sort all points by that distance.
  • Return the first k.

Why it is weak:

  • Sorting all n points costs O(n log n).
  • When n is huge and k is small, that is a lot of extra work.
  • You sorted points you throw away anyway.

Here is the sort-all-points code:

k_closest_points_sorting.py
def k_closest(points, k):
    points.sort(key=lambda p: p[0] * p[0] + p[1] * p[1])
    return points[:k]

⚡ Approach 2: A Max-Heap of Size k (Best)

The idea in one line: keep only k points and drop the farthest one fast.

What a max-heap is:

  • A max-heap is a binary tree kept in an array with the largest value at the top.
  • Reading the top is instant. Adding or removing is O(log k) when it holds k items.
  • It is also called a priority queue.

How it works:

  • Order the heap by distance, so the farthest point sits at the top.
  • Walk through all points and push each one.
  • If the heap grows past k, pop the top. The top is the farthest point.
  • When you finish, the heap holds exactly the k closest points.

Why a max-heap, not a min-heap:

  • We want to throw away the farthest point quickly.
  • The max-heap keeps that farthest point at the top, ready to remove.

Why it is fast:

  • Each push or pop on a heap of size k is O(log k).
  • You only ever store k points, so extra space is O(k).

Here is the max-heap of size 2 drawn as a tree after we process all points. The farther of the two kept points, distance 8, sits at the top.

dist2 = 8 -> (-2,2)

dist2 = 1 -> (0,1)

Steps to Solve

  1. Make an empty max-heap ordered by distance. It holds at most k points.
  2. For each point, compute its distance squared, which is x*x + y*y.
  3. Push the point onto the heap using that distance.
  4. If the heap has more than k points, pop the top. The top is the farthest point.
  5. After all points, the heap holds the k closest points. Return them.

Python’s heapq is a min-heap. To make a max-heap on distance, we store the negative distance. The most negative, which is the farthest point, then sits on top.

k_closest.py
import heapq




def k_closest(points, k):
    heap = []  # max-heap on distance, using negative distance
    for x, y in points:
        dist = x * x + y * y                 # distance squared, no square root
        heapq.heappush(heap, (-dist, x, y))  # negate so farthest is on top
        if len(heap) > k:
            heapq.heappop(heap)              # drop the farthest point
    return [[x, y] for (_, x, y) in heap]




points = [[1, 3], [-2, 2], [5, 8], [0, 1]]
for p in sorted(k_closest(points, 2), key=lambda p: p[0]*p[0] + p[1]*p[1]):
    print(p)

The output of the above code will be:

[0, 1]
[-2, 2]

Let us walk through the Python version line by line. The max-heap on negative distance does the heavy lifting.

We start with an empty list called heap. For each point we read its x and y. The line dist = x * x + y * y is the distance squared. We skip the square root, because comparing the squares gives the same order. That saves a math call per point.

The push line stores a tuple (-dist, x, y). We negate the distance. So the most negative value, which is the farthest point, ends up on top of the heap. Python compares tuples by the first item first, so the heap orders by distance.

The if len(heap) > k checks if we are holding too many points. If so, heapq.heappop removes the top. The top is the farthest point. So we always drop the worst and keep the k closest.

After the loop, the heap holds exactly k points. The last line builds a clean list of [x, y] pairs, dropping the distance we no longer need. The sort in the print step is only there to show the output in a steady order. The heap itself already holds the right k points.

⏱️ Time and Space Complexity

The full sort is O(n log n). The heap of size k walks through all n points, and each push or pop is O(log k). So the heap approach is O(n log k). When k is much smaller than n, that is a real saving. The heap holds at most k points, so its extra space is O(k). The max-heap lets us keep only what we need.

Approach Time Complexity Space Complexity
Sort all points O(n log n) O(n)
Max-heap of size k O(n log k) O(k)

Tip

Use distance squared, not the real distance. The square root does not change which point is closer, so skip it and save a calculation on every point.

🧩 Key Takeaways

  • ✅ Compare distance squared, which is x*x + y*y, and skip the square root.
  • ✅ A max-heap of size k keeps the farthest of the k at the top, ready to drop.
  • ✅ Walk through all points, push each, and pop when the heap grows past k.
  • ✅ The heap approach is O(n log k), faster than a full sort when k is small.
  • ✅ You only store k points, so extra space stays at O(k).

Check Your Knowledge

 4 questions Show quiz Hide quiz

Test what you learned. Pick an answer for each question, then click Check.

  1. 1

    Why do we compare x*x + y*y instead of the real distance?

    Why: Distance squared keeps the same order as the real distance, so we skip the square root and save a calculation.

  2. 2

    Why use a max-heap and not a min-heap to keep the k closest points?

    Why: We want to throw away the farthest point fast, and the max-heap keeps that farthest point ready at the top.

  3. 3

    What is the time complexity of the max-heap of size k approach?

    Why: We process n points, and each push or pop on a heap of size k is O(log k), giving O(n log k).

  4. 4

    How many points does the heap hold at most?

    Why: Whenever the heap grows past k, we pop the farthest, so it never holds more than k points.

Score: 0 / 4

🚀 What’s Next?

Previous Last Stone Weight Next Task Scheduler