Monday, October 2, 2017

K closest points

/*
 * Find the K closest points to the origin in 2D plane, given an array containing N points.
 * You can assume K is much smaller than N and N is very large.
 * You need only use standard math operators (addition, subtraction, multiplication, and division).
 */
import java.util.*;

class Point {
    double x;
    double y;
    public Point(double x, double y) {
        this.x = x;
        this.y = y;
    }
}
public class Solution {
private static double distance(Point a, Point b) {
        return (a.x - b.x) * (a.x - b.x) + (a.y - b.y) * (a.y - b.y);
    }
public static Point[] closestPoint(Point[] array, final Point origin, int k) {
        if (k >= array.length) {
        // return array;
        k = array.length;
        }
        if (k <= 0) {
        return new Point[0];
        }
        Point[] res = new Point[k];
        PriorityQueue<Point> maxHeap = new PriorityQueue<>(k + 1, new Comparator<Point>() {
            @Override
            public int compare(Point a, Point b) {
                return Double.compare(distance(b, origin), distance(a, origin));
            }
        });
        for(Point p : array) {
        maxHeap.offer(p);
        if (maxHeap.size() > k) {
        maxHeap.poll();
        }
        }
        for(int i = k - 1; i >= 0; i--) {
        res[i] = maxHeap.poll();
        }
        return res;
    }

    public static void main(String[] args) {
        Point origin = new Point(0, 0);
        Point[] input = new Point[]{new Point(0, 2), new Point(1, 1), new Point(-1, 0), new Point(2, 0), new Point(3, 0)};
        Point[] output = closestPoint(input, origin, 3);
        System.out.println("input");
        for(Point i : input) System.out.print("("+i.x+", "+i.y+") ");
        System.out.println("");
        System.out.println("output");
        for(Point i : output) System.out.print("("+i.x+", "+i.y+") ");
    }
}

No comments:

Post a Comment