Java—Leetcode 973. K Closest Points to Origin

Description

Given an array of points where points[i] = [xi, yi] represents a point on the X-Y plane and an integer k, return the k closest points to the origin (0, 0).

The distance between two points on the X-Y plane is the Euclidean distance (i.e., √(x1 - x2)2 + (y1 - y2)2).

You may return the answer in any order. The answer is guaranteed to be unique (except for the order that it is in).

Code

public int[][] kClosest(int[][] points, int k) {
        PriorityQueue<int[]> priorityQueue = new PriorityQueue<>(new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                return (o2[0] * o2[0] + o2[1] * o2[1]) - (o1[0] * o1[0] + o1[1] * o1[1]);
            }
        }
        );

        for (int[] point : points) {
            priorityQueue.offer(point);
            if (priorityQueue.size() > k)
                priorityQueue.poll();
        }

        int[][] res = new int[k][2];
        int i = 0;
        while (!priorityQueue.isEmpty())
            res[i++] = priorityQueue.poll();

        return res;
    }