K Closest Points To Origin (Interview Question and Solution)

How to Solve K Closest Points To Origin

K Closest Points To Origin Introduction

The K Closest Points To Origin problem involves comparing the distance of points plotted on a graph. This is a common problem in data analysis, most often found in the context of generating nearest neighbor sets. Similar to other k-selection algorithms, this problem can be solved with a variety of sorting techniques and challenges us to use a heap data structure to improve time complexity. Before viewing the problem and solution, below are some short video snippets from real mock interviews to help prepare you for some common pitfalls that interviewees stumble into.

• Problem

• Solution

• Interview Analysis: Snippets from Real Interviews 🔥

K Closest Points To Origin Problem

Given a list of tuples that represent (X, Y) coordinates on an XY plane and an integer K, return a list of the K-closest points to the origin (0, 0).

Example Inputs and Outputs

Example 1

Input:

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

Output:

[[3, 3], [4, 4]] or [[4, 4], [3, 3]]

Example 2

Input:

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

Output:

[[-1, -1], [1, 2]] or [[1, 2], [-1, -1]]

Constraints

The number of nodes in the list is in the range [0, 5000]

K is >= 0 and <= the length of the input list

K Closest Points To Origin Solutions

To solve this problem we will need to do some basic algebra. We have a right triangle and we need to calculate the length of the hypotenuse, and we can do so using the Pythagorean theorem:

A^2 + B^2 = C^2

As a quick aside, we don't actually need to calculate C (the hypotenuse, or distance from origin), as simply calculating A^2 + B^2 for each coordinate will allow us to order the points from closest to furthest from origin without actually determining the exact distance.

1. Sorting

Given we need to find the K closest points to origin, the naive approach should hopefully become clear relatively quickly. If we calculate the distance for each coordinate pair, we can then sort the coordinates by distance and finally slice the list from 0 to K in order to return the K closest points to the origin.

K Closest Points to Origin Python and JavaScript Solutions - Sorting


def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
  points.sort(key=lambda xy_tuple: xy_tuple[0]**2 + xy_tuple[1]**2)
  return points[:k]

Time/Space Complexity Analysis

Time Complexity: O(n * log(n))
Space Complexity: O(k), as we are sorting in place and returning a new list with K points

2. Using a Heap

To improve our time complexity we will need to avoid fully sorting the input, and if we aren't sorting the input we will need to repeatedly select the point with the smallest distance from the origin.

To give a quick refresher, heaps are an ordered (but not fully sorted) data structure often backed by an array. They can be created in linear time and they ensure selection of the smallest or largest element at any given time.

K Closest Points to Origin Python and JavaScript Solutions - Using a Heap


import heapq

def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
  distance_coordinate_tuples = [(x*x + y*y, [x, y]) for x, y in points]
  heapq.heapify(distance_coordinate_tuples)
  k_smallest = heapq.nsmallest(k, distance_coordinate_tuples)
  return [coordinate for distance, coordinate in k_smallest]

Time/Space Complexity Analysis

Time Complexity: O(n) + O(k*log(n))
Space Complexity: O(1) (or O(n) if you mutate the input)

K Closest Points To Origin Analysis

Common Mistakes: Forgetting Data Structures

Often candidates may have a great understanding of programming and problem solving skills. However if you haven't practised in a while, you may find you have forgotten what some data structures look like or how to fully implement them. Practise makes perfect, so ensure you brush up before you interview!

Common Mistakes: Syntax Errors

Prospective candidates will often be knowledgeable in more than one programming language. During mock technical interviews, it is important to practise and familiarize yourself with the language you will be interviewing in, especially if it differs from the native language you code in at your current job.

Common Mistakes: Variable Names and Data Types

When calculating distance in this question, it is important to use the double data type as opposed to int. This snippet shows the importance of looking out for small mistakes in variable data types and names that can lead to errors when running code.

Senior Level Extension Question

Many solutions to the K closest points question do not take into account what would happen if given a large data set of points or a near infinite stream of points. A MapReduce is one solution to this problem.

Possible Approach: Checking Distance

In this snippet the interviewer suggests checking if it's the first k instead of checking if the distance is less than the minimum.

Optimization: Size Efficiency

In this snippet the candidate suggests optimizing the comparator function by using a map to store distance-to-pair values, rather than recalculating distances each time.