K Closest Points to Origin

Given a list of points on the plane.  Find the K closest points to the origin (0, 0). The distance between two points on a plane is the Euclidean Distance.

Example 1:

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

Explanation:

A point (x',y') is given in a 2D axis. So we can form a rectangle by using the origin (0,0) and the given point (x',y'). Now the diagonal of this rectangle is the distance of the given point (x',y') from the origin (0,0). We can use Pythagoras’ Theorem to find the length of the diagonal if we know the width and height of the rectangle. As a formula:

Diagonal = SQRT (h2 + w2)
Where:
w  is the width of the rectangle
h  is the height of the rectangle

Now use this formula for the given example. 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]]

Note:

  1. 1 <= K <= points.length <= 10000
  2. -10000 < points[i][0] < 10000
  3. -10000 < points[i][1] < 10000

Consider the following figure for a better understanding:

The very naive and simple solution is sorting all points by their distance to the origin point directly, then get the top k closest points. We can use the sort function and the code is very short. The time complexity of this approach is O(NlogN), due to the sorting. The advantages of this solution are short, intuitive and easy to implement. The disadvantage of this solution is not very efficient and have to know all of the points previously. With this approach, you can not handle a stream of continuous points as your input. The implementation is as follows:

import java.util.*;

class Point implements Comparable<Point> {
    int x;
    int y;

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

    public int compareTo(Point point) {
        return ((x * x) + (y * y)) - ((point.x * point.x) + (point.y * point.y));
    }
}

class KClosestPointToOriginSol {
    public int[][] kClosest(int[][] points, int K) {
        List<Point> pointList = new ArrayList<>();
        for (int[] point : points) {
            Point p = new Point(point[0], point[1]);
            pointList.add(p);
        }
        Collections.sort(pointList);
        int[][] ans = new int[K][2];
        for (int i = 0; i < K; i++) {
            Point point = pointList.get(i);
            ans[i][0] = point.x;
            ans[i][1] = point.y;
        }
        return ans;
    }
}

public class KClosestPointToOrigin {

    public static void main(String[] args) {
        int[][] points = {{3, 3}, {5, -1}, {-2, 4}};
        int k = 2;
        KClosestPointToOriginSol solution = new KClosestPointToOriginSol();
        int[][] result = solution.kClosest(points, k);
        System.out.println("K Closest points are");
        for (int[] point : result) {
            System.out.println("[x:" + point[0] + ",y:" + point[1] + "]");
        }
    }
}

The output of the program is:

K Closest points are [x:3,y:3] [x:-2,y:4]

Use a Max Heap to Handle Real-Time input:

The second solution is based on the first one. We don’t have to sort all points. Instead, we can maintain a MaxHeap with size K. Then for each point, we add it to the heap. Once the size of the heap is greater than K, we are supposed to extract one from the max-heap to ensure the size of the heap is always K. Thus, the max-heap is always maintained top K smallest elements from the first one to the current one. Once the size of the heap is over its maximum capacity, it will exclude the maximum element in it, since it can not be the proper candidate anymore. The time complexity is O(NlogK)The advantage of this solution is it can deal with real-time(online) stream data. It does not need to know the size of the data previously. Check the following implementation:

import java.util.Comparator;
import java.util.PriorityQueue;

class Point implements Comparable<Point> {
    int x;
    int y;

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

    public int compareTo(Point point) {
        return ((x * x) + (y * y)) - ((point.x * point.x) + (point.y * point.y));
    }
}

class MaxHeap extends PriorityQueue<Point> {

    public MaxHeap(int size) {
        super(size, new Comparator<Point>() {
            public int compare(Point o1, Point o2) {
                return o2.compareTo(o1);
            }
        });
    }
}

class Solution_2 {
    public int[][] kClosest(int[][] points, int K) {
        MaxHeap maxHeap = new MaxHeap(K);
        for (int[] point : points) {
            Point p = new Point(point[0], point[1]);
            maxHeap.add(p);
            if (maxHeap.size() > K) {
                maxHeap.poll();
            }
        }
        int i = 0;
        int[][] result = new int[K][2];
        for (Point point : maxHeap) {
            result[i][0] = point.x;
            result[i][1] = point.y;
            i++;
        }
        return result;
    }
}

public class KClosestPointToOrigin {
    public static void main(String[] args) {
        int[][] points = {{3, 3}, {5, -1}, {-2, 4}};
        int k = 2;
        Solution_2 solution = new Solution_2();
        int[][] result = solution.kClosest(points, k);
        System.out.println("K Closest points are");
        for (int[] point : result) {
            System.out.println("[x:" + point[0] + ",y:" + point[1] + "]");
        }
    }
}

The output of the program is:

K Closest points are [x:-2,y:4] [x:3,y:3]

One improvement that can be done in the above solution is, rather than adding every point to the heap we can do the following:

  1. Add every point to the heap if the heap size is less than K.
  2. When heap size is K, then add the current point in the heap only if it is closer to the origin than the Heap top.
  3. If the heap top is far from the origin compared to the incoming point then add the incoming point in the heap and remove the earlier top.
  4. Although this improvement will not change the runtime complexity, which is still O(NlogK).
 public int[][] kClosest(int[][] points, int K) {
        MaxHeap maxHeap = new MaxHeap(K);
        for (int[] point : points) {
            Point currentPoint = new Point(point[0], point[1]);
            if (maxHeap.size() < K) {
                maxHeap.add(currentPoint);
            } else {
                Point top = maxHeap.peek();
                if (currentPoint.compareTo(top) < 0) {
                    maxHeap.add(currentPoint);
                    maxHeap.remove();
                }
            }
        }
        int[][] result = new int[K][2];
        int i = 0;
        for (Point point : maxHeap) {
            result[i][0] = point.x;
            result[i][1] = point.y;
            i++;
        }
        return result;
}

Categories: Heap

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s