Problem

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).

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.

Brainstorming

okay so this looks like we need 2 things ,

1. a way to calculate the distance between a point and the origin
2. a data structure to store the result and and the point

we do this for all the points , then we return the k nearest point

A min heap works the best here this is similar to all the problems we have solved in heaps before like [[Top K Frequent Words]] [[Kth Largest Element]] etc

1. we define a helper which calculates the distance of a point from the origin using the euclidian formula given
2. we define a min heap
3. for each point in the array we iterate and get the distance
4. the length of min heap is < K we just add the packet i.e tuple ( distance , (point))
5. otherwise we check if the current distance packet is smaller than the smallest distance packet till now in the heap, if yes we just replace it with the current distance packet
6. at the end of the loop we will have k nearest points to the origin

import heapq
import math

points = [[3,3],[5,-1],[-2,4]]
k = 2

def euclidian_distance_helper(point):
	origin = [0,0]
	distance_squared = (origin[0]-point[0])**2 + (origin[1]-point[1])**2
	distance_square_root = math.sqrt(distance_squared)
	return distance_square_root

def k_nearest_point(points, K):
	min_heap = []
	results = []
	for point in points:
		distance = euclidian_distance_helper(point)
		if len(min_heap) < K:
			heapq.heappush(min_heap, (distance,point))
		else:
			current_nearest_point_packet = min_heap[0]
			if distance < current_nearest_point_packet[0]:
				heapq.heappushpop(min_heap, (distance, point))
		
	while min_heap:
		current_nearest_point_packet = heapq.heappop(min_heap)
		results.append(current_nearest_point_packet[1])
	
	return results

print(k_nearest_point(points,k))

• 🍅 (pomodoro::WORK) (duration:: 20m) (begin:: 2025-12-14 09:46) - (end:: 2025-12-14 10:09)

Optimal solution

Let’s try points = [[3,3], [5,-1], [-2,4]]. K=2.
Distances:

- [3,3]: sqrt(18) ~ 4.24 - [5,-1]: sqrt(26) ~ 5.1

- [-2,4]: sqrt(20) ~ 4.47

Correct answer should be [[3,3], [-2,4]].

Let’s trace your code again:

1. Heap is [].

2. Process [3,3] (dist 4.24). heappush. Heap: [(4.24, [3,3])].

3. Process [5,-1] (dist 5.1). heappush. Heap: [(4.24, [3,3]), (5.1, [5,-1])].

4. Process [-2,4] (dist 4.47). len is now 2. Go to else.

   • current_nearest is (4.24, [3,3]) (the smallest).
   • if 4.47 < 4.24 is false. Nothing happens.

5. Loop ends. Heap contains [(4.24, [3,3]), (5.1, [5,-1])]. Your code will return [[3,3], [5,-1]]. This is incorrect.

The Crux of the Issue:
You are trying to keep a collection of the K smallest items seen so far. When a new item arrives, you need to compare it to the LARGEST of the items in your collection. If the new item is smaller than the largest one you’re currently holding, you kick out the largest and insert the new, smaller one.

A min-heap is very good at giving you the smallest item in O(1). It is very inefficient at giving you the largest item.

Question (The Self-Nudge):

1. What kind of heap is extremely efficient at finding and removing the largest element?
2. How would you simulate this type of heap using Python’s heapq, which is a min-heap by default? (Hint: what happens if you store negative distances in the heap?)

(Prompt): Answer these two questions. This is the key to flipping your algorithm to the correct, optimal pattern for a “Top K Smallest” problem.

we need to use a max heap

but in python we only have min heap by default , but if we flip the sign of the distance then it behaves like a max heap

import heapq
import math

points = [[3,3],[5,-1],[-2,4]]
k = 2

def euclidian_distance_helper(point):
	origin = [0,0]
	distance_squared = (origin[0]-point[0])**2 + (origin[1]-point[1])**2
	distance_square_root = math.sqrt(distance_squared)
	return distance_square_root

def k_nearest_point(points, K):
	max_heap = []
	results = []
	for point in points:
		distance = euclidian_distance_helper(point)
		if len(max_heap) < K:
			heapq.heappush(max_heap, (-distance,point))
		else:
			current_nearest_point_packet = max_heap[0]
			if - distance > current_nearest_point_packet[0]:
				heapq.heappushpop(max_heap, (-distance, point))
		
	while max_heap:
		current_nearest_point_packet = heapq.heappop(max_heap)
		results.append(current_nearest_point_packet[1])
	
	return results

print(k_nearest_point(points,k))

Final Polish (Optional Optimizations):

1. Tie-breaker: Just like in the mergeKLists problem, if two points had the exact same distance, Python might try to compare the point lists, which can be slow or unpredictable. The robust solution is to add a unique ID, but for LeetCode, this is often not necessary.
2. Square Roots: The sqrt calculation is relatively slow. Do you actually need it? Comparing d1 and d2 is the same as comparing d1d1 and d2d2 (as long as distances are non-negative). You can optimize the code by working with the squared distances directly and avoiding the math.sqrt call entirely. This is a common and valuable optimization to mention in an interview.

But as it stands, your code is a complete and correct O(N log K) solution, which is the optimal time complexity for this problem. You have successfully identified and implemented the “Top K” pattern using a max-heap. Fantastic work.