K Closest Points to the Origin

Innoskrit August 18, 2022 Innoskrit Masterclass DSA Sheet, Priority Queue (Heap)

Problem Statement

Given an array of points in a 2D plane, find ‘K’ closest points to the origin.

Example:

Input 1: points = [[1,2],[1,3]], K = 1
Output 1: [[1,2]]
Explanation 1: The Euclidean distance between (1, 2) and the origin is sqrt(5). The Euclidean distance between (1, 3) and the origin is sqrt(10).
Since sqrt(5) < sqrt(10), therefore (1, 2) is closer to the origin.

Input 2: point = [[1, 3], [3, 4], [2, -1]], K = 2
Output 2: [[1, 3], [2, -1]]

Code

Java

C++

Python

import java.util.*;

class Point {
  int x;
  int y;

  public Point(int x, int y) {
    this.x = x;
    this.y = y;
  }

  public int distFromOrigin() {
    return (x * x) + (y * y);
  }
}

class Innoskrit {

  public static List<Point> findClosestPoints(Point[] points, int k) {
    PriorityQueue<Point> maxHeap = new PriorityQueue<>((p1, p2) -> p2.distFromOrigin() - p1.distFromOrigin());

    for (int i = 0; i < k; i++)
      maxHeap.add(points[i]);

    for (int i = k; i < points.length; i++) {
      if (points[i].distFromOrigin() < maxHeap.peek().distFromOrigin()) {
        maxHeap.poll();
        maxHeap.add(points[i]);
      }
    }

    return new ArrayList<>(maxHeap);
  }

  public static void main(String[] args) {
    Point[] points = new Point[] { new Point(1, 3), new Point(3, 4), new Point(2, -1) };
    List<Point> result = Innoskrit.findClosestPoints(points, 2);
    for (Point p : result)
      System.out.print("[" + p.x + " , " + p.y + "] ");
  }
}

using namespace std;

#include <iostream>
#include <queue>
#include <vector>

class Point {
 public:
  int x = 0;
  int y = 0;

  Point(int x, int y) {
    this->x = x;
    this->y = y;
  }

  int distFromOrigin() const {
    return (x * x) + (y * y);
  }

  const bool operator<(const Point& p) { return p.distFromOrigin() > this->distFromOrigin(); }
};

class Innoskrit {
 public:
  static vector<Point> findClosestPoints(const vector<Point>& points, int k) {

    vector<Point> maxHeap(points.begin(), points.begin() + k);
    make_heap(maxHeap.begin(), maxHeap.end());

    for (int i = k; i < points.size(); i++) {
      if (points[i].distFromOrigin() < maxHeap.front().distFromOrigin()) {
        pop_heap(maxHeap.begin(), maxHeap.end());
        maxHeap.pop_back();
        maxHeap.push_back(points[i]);
        push_heap(maxHeap.begin(), maxHeap.end());
      }
    }

    return maxHeap;
  }
};

int main(int argc, char* argv[]) {
  vector<Point> maxHeap = Innoskrit::findClosestPoints({{1, 3}, {3, 4}, {2, -1}}, 2);
  for (auto p : maxHeap) {
    cout << "[" << p.x << " , " << p.y << "] ";
  }
}

from heapq import *

class Point:

  def __init__(self, x, y):
    self.x = x
    self.y = y

  def __lt__(self, other):
    return self.distance_from_origin() > other.distance_from_origin()

  def distance_from_origin(self):
    return (self.x * self.x) + (self.y * self.y)

  def print_point(self):
    print("[" + str(self.x) + ", " + str(self.y) + "] ", end='')

def find_closest_points(points, k):
  maxHeap = []
  for i in range(k):
    heappush(maxHeap, points[i])

  for i in range(k, len(points)):
    if points[i].distance_from_origin() < maxHeap[0].distance_from_origin():
      heappop(maxHeap)
      heappush(maxHeap, points[i])

  return maxHeap

def main():

  result = find_closest_points([Point(1, 3), Point(3, 4), Point(2, -1)], 2)
  for point in result:
    point.print_point()

main()

Output:

[1, 3], [2, -1]

Time and Space Complexity

Time Complexity: The time complexity of this algorithm is O(N*logK).

Space Complexity: The space complexity will be O(K) since we need to store the top ‘K’ points in the heap.