Start learning free

LESSON 385 OF 536 · HEAP

K Closest Points to Origin in Java: O(n log k) with a Max-Heap

Practice O(n log k) Read the full explanation

Given an array of points, where points[i] = [x, y] is a point on the plane, and an integer k, return the k points closest to the origin (0, 0).

The distance between a point [x, y] and the origin is the usual straight-line distance, √(x2 + y2).

Return the points sorted by their distance to the origin, the closest first. Points at the same distance come in order of their x, and points with the same x too in order of their y. The answer is guaranteed to be unique: every point left out is farther from the origin than every point in the answer.

 

Example 1:

Input: points = [[1,3],[-2,2]], k = 1
Output: [[-2,2]]
Explanation: [1,3] is √10 away from the origin and [-2,2] is √8 away.
√8 is smaller, so [-2,2] is the closest point.

Example 2:

Input: points = [[3,3],[5,-1],[-2,4]], k = 2
Output: [[3,3],[-2,4]]
Explanation: The squared distances are 18, 26 and 20. The two closest points are [3,3] (18) and [-2,4] (20), in that order.

Example 3:

Input: points = [[0,2],[2,0],[1,1],[0,-2],[3,3]], k = 4
Output: [[1,1],[0,-2],[0,2],[2,0]]
Explanation: [1,1] is the closest (√2). [0,2], [2,0] and [0,-2] are all 2 away, so they come by x:
[0,-2] and [0,2] (x = 0) before [2,0], and [0,-2] before [0,2] because -2 < 2. [3,3] is the farthest and is left out.

 

Constraints:

  • 1 <= k <= points.length <= 105
  • -104 <= x, y <= 104
  • The same point can appear more than once.
  • The answer is unique.

Note:

Your algorithm should run in O(n log k) time and use O(k) extra space, where n is the number of points. Sorting all the points by distance would take O(n log n).

No square roots are needed: one point is closer than another exactly when its x2 + y2 is smaller, and that squared distance is a whole number of at most 2 · 108, which fits in a 32-bit integer.

STEP 1 OF 12
    Ctrl + Enter to check

    JAVA · HEAP

    K Closest Points To Origin in Java, explained

    Problem overview

    The method kClosest receives points, where points[i] = [x, y] is a point on the plane, and an integer k. It returns the k points closest to the origin (0, 0). The distance is the usual straight-line one, the square root of x * x + y * y.

    The answer is sorted: closest first, ties broken by x, then by y. Every point left out is farther away than every point in the answer. Some cases:

    • [[1,3],[-2,2]] with k = 1 gives [[-2,2]], because the squared distances are 10 and 8.
    • [[3,3],[5,-1],[-2,4]] with k = 2 gives [[3,3],[-2,4]]. The squared distances are 18, 26 and 20.
    • [[0,2],[2,0],[1,1],[0,-2],[3,3]] with k = 4 gives [[1,1],[0,-2],[0,2],[2,0]]. Three points are at squared distance 4, so they come ordered by x, then y.

    The statement asks for O(n log k) time and O(k) extra space, where n is the number of points. Sorting all the points by distance would cost O(n log n).

    Approach: a max-heap of size k

    Two ideas combine here.

    No square roots. One point is closer than another exactly when its x * x + y * y is smaller, so compare squared distances. They are whole numbers. With coordinates up to 10,000 in absolute value the squared distance is at most 200,000,000 (2 * 10000 * 10000), which fits easily in an int.

    Keep only the best k. Hold the k closest points seen so far in a heap with the farthest of them on top, which is a max-heap on the distance. When a new point arrives, add it. If the heap now has more than k points, the one on top is the farthest of k + 1 points, so it can't be among the k closest, and it is removed. After the last point, the heap holds exactly the answer. Each add or removal costs O(log k), so the whole pass is O(n log k), and the heap never holds more than k + 1 entries.

    Take the third example with k = 4. Adding the points one by one gives:

    added [0,2] dist 4, size 1, top dist 4
added [2,0] dist 4, size 2, top dist 4
added [1,1] dist 2, size 3, top dist 4
added [0,-2] dist 4, size 4, top dist 4
added [3,3] dist 18, size 5, top dist 18, removed [3,3]

    The last point is the farthest, so it is added and immediately removed. A heap is only partly ordered, though: it guarantees the top and nothing about the rest. That is why the lesson sorts the survivors before returning them.

    Step by step

    1. Create the heap: a PriorityQueue<int[]> called heap with the comparator (a, b) -> b[0] - a[0]. Each entry is an array {dist, x, y}, and this comparator puts the biggest distance on top. The subtraction can't overflow because distances are at most 200,000,000.
    2. Loop over every point with a for-each loop over points.
    3. Read the coordinates: int x = point[0], y = point[1].
    4. Compute dist, the squared distance x * x + y * y.
    5. Add the entry new int[]{dist, x, y} to the heap. It keeps the point together with the number the heap orders by.
    6. Trim the heap. If heap.size() > k, call heap.poll() to remove the top, the farthest point.
    7. Copy the survivors into an int[][] named closest with heap.toArray(new int[0][]).
    8. Sort closest with Arrays.sort and a comparator that compares the distances at index 0. When those are equal it compares x at index 1, and when those are equal too it compares y at index 2. The cases are chained with ? :.
    9. Build the answer: create result as new int[k][], with k rows and no columns allocated yet.
    10. Fill it in a loop with i from 0 to k - 1, setting result[i] to new int[]{closest[i][1], closest[i][2]}. The distance at index 0 has done its job, so only the coordinates go in.
    11. Return result.

    The solution in Java

    class Solution {
    public int[][] kClosest(int[][] points, int k) {
        PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> b[0] - a[0]);
        for (int[] point : points) {
            int x = point[0], y = point[1];
            int dist = x * x + y * y;
            heap.add(new int[]{dist, x, y});
            if (heap.size() > k) {
                heap.poll();
            }
        }
        int[][] closest = heap.toArray(new int[0][]);
        Arrays.sort(closest, (a, b) -> a[0] != b[0] ? a[0] - b[0] : a[1] != b[1] ? a[1] - b[1] : a[2] - b[2]);
        int[][] result = new int[k][];
        for (int i = 0; i < k; i++) {
            result[i] = new int[]{closest[i][1], closest[i][2]};
        }
        return result;
    }
}

    Java details worth knowing:

    • PriorityQueue is a min-heap by default. The comparator (a, b) -> b[0] - a[0] reverses that, so the largest distance is at the head, and poll() removes it.
    • Comparator by subtraction. A comparator returns a negative number, zero or a positive number, and a - b does that for small values. It is only safe when the difference can't overflow. For example, Integer.MIN_VALUE - 1 wraps around to 2147483647. Here the values are small, so it is fine.
    • The heap's own order is partial. Printing a PriorityQueue built from 5, 1, 4, 2, 3 shows [1, 2, 4, 5, 3], which is not sorted. toArray returns the elements in that internal order, so the explicit sort is needed.
    • toArray(new int[0][]). The empty array tells Java which array type to build; the method returns a new int[][] of the right size.
    • new int[k][]. This creates k rows that are all null until assigned. The loop fills each one, and rows may be created later because Java arrays of arrays don't need equal row lengths.
    • The nested ? :. It reads left to right: if the distances differ, use their difference; otherwise, if the x values differ, use that difference; otherwise use the y difference.

    Complexity

    Time: O(n log k). Each of the n points costs one add and at most one poll on a heap of at most k + 1 entries, which is O(log k) each. The final sort handles k entries and costs O(k log k), which is within the bound because k is at most n.

    Space: O(k). The heap, the closest array and the result array each hold at most k small arrays.

    Edge cases and common mistakes

    • One point. [[5,5]] with k = 1 gives [[5,5]].
    • Duplicate points. The same point can appear more than once. [[1,1],[1,1],[2,2]] with k = 2 gives [[1,1],[1,1]].
    • Points at equal distance. [[10000,10000],[-10000,-10000]] with k = 2 gives [[-10000,-10000],[10000,10000]]: equal distances, so the smaller x comes first.
    • Using a min-heap. With the default order the closest point sits on top, so the wrong point is removed when the heap grows past k.
    • Removing before adding. The new point could be closer than anything in the heap, so add it first and then drop the farthest one.
    • Taking a square root. It is not needed, and comparing exact integers avoids floating-point rounding.
    • Returning the heap's order. Without the sort the points come out in the heap's internal order, which fails the sorted-output requirement.

    Testing

    • points = [[1,3],[-2,2]], k = 1 gives [[-2,2]].
    • points = [[3,3],[5,-1],[-2,4]], k = 2 gives [[3,3],[-2,4]].
    • points = [[0,2],[2,0],[1,1],[0,-2],[3,3]], k = 4 gives [[1,1],[0,-2],[0,2],[2,0]].
    • points = [[1,1],[1,1],[2,2]], k = 2 gives [[1,1],[1,1]].
    Keep going with Java500+ step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free