Katkat
Back

Find the K Closest Points to a Query Location

UberSoftware Engineer · Onsite
MediumSoftware EngineerOnsiteCoding & Algorithms

Last updated

This problem tests the ability to efficiently find the k nearest neighbors using a priority queue (max-heap) and custom sorting logic for tie-breaking.

Problem Statement

Given a list of points on a 2D plane represented as [x, y] coordinates, a query location [qx, qy], and an integer k, return the k closest points to the query location.

Distance should be calculated using the squared Euclidean distance: (x - qx)^2 + (y - qy)^2. Do not use square roots to avoid floating-point precision issues.

Tie-breaking Rules

If two points have the same distance to the query location, use the following rules to determine the order:

  1. If distances are equal, the point with the smaller x-coordinate is considered "closer."
  2. If x-coordinates are also equal, the point with the smaller y-coordinate is considered "closer."

Function Signature

def find_k_closest(points: List[List[int]], query: List[int], k: int) -> List[List[int]]:

Example 1

Input: points = [[1, 3], [-2, 2], [5, 8], [0, 1]], query = [0, 0], k = 2 Output: [[0, 1], [-2, 2]] Explanation:

  • Distances squared: [1, 3] -> 10, [-2, 2] -> 8, [5, 8] -> 89, [0, 1] -> 1
  • Sorted order: [0, 1] (dist 1), [-2, 2] (dist 8), [1, 3] (dist 10), [5, 8] (dist 89).
  • Top 2 are [0, 1] and [-2, 2].

Constraints

  • 1 <= k <= len(points)
  • -10^4 <= x, y <= 10^4
  • The number of points n can be up to 10^5.

Hints

Hint 1

Instead of sorting the entire list, consider using a data structure that keeps track of the 'top k' elements seen so far.

Hint 2

A max-heap can store the current k closest points; if a new point is closer than the heap's maximum, replace the maximum.

Hint 3

Python's `heapq` module creates a min-heap by default; to implement a max-heap, store tuples with negated values.

Follow-up questions

  • How would you modify your solution if the points were stored on disk and couldn't fit into memory?
  • If the query location changes frequently but the set of points remains static, what indexing structure could you build to improve performance?
  • How does your solution perform if the coordinates are in 3D space instead of 2D?