Problem Overview

You are given an array of points where points[i] = [xi, yi] represents a point on the X-Y plane, and an integer k.

Your task is to return the k points closest to the origin (0, 0).

The distance between any point (x, y) and the origin (0, 0) is calculated as:

distance = √(x² + y²)

However, since we only need to compare distances, we can use the squared distance x² + y² to avoid unnecessary square root calculations.

Example 1:
Input:
points = [[1,3], [-2,2]], k = 1
Output:
[[-2,2]]
Explanation:
The distance of each point from the origin is:

• [1,3] → √(1² + 3²) = √10
• [-2,2] → √(4 + 4) = √8

Since √8 < √10, the closest point is [-2, 2].

Example 2:
Input:
points = [[3,3],[5,-1],[-2,4]], k = 2
Output:
[[3,3],[-2,4]]
Explanation:
The distances are:

• [3,3] → 18
• [5,-1] → 26
• [-2,4] → 20

The two smallest distances are 18 and 20, corresponding to [3,3] and [-2,4].

---

Intuition / Approach

The problem asks for the k smallest distances from the origin among the given points. This can be viewed as a partial sorting problem.

There are several possible approaches:

1. Sort All Points by Distance:
   • Compute the distance for each point.
   • Sort all points based on their squared distance.
   • Pick the first k points.
     This is simple but less efficient for large datasets, as sorting takes O(n log n) time.
2. Use a Max-Heap of Size K:
   • Maintain a heap (priority queue) that stores the k closest points seen so far.
   • For each new point, calculate its distance:
     • If the heap has fewer than k points, add it.
     • Otherwise, compare the distance of the new point with the maximum distance in the heap:
       • If smaller, remove the farthest point and add the new one.
   • This ensures the heap always contains the k closest points.
     This approach is efficient when n is large and k is much smaller than n.

We will use the max-heap approach, as it’s more scalable and optimal in most practical scenarios.

---

Algorithm Explanation

1. Create a priority queue (max-heap) where the point with the largest distance is at the top.
2. Iterate over all points:
   • Compute the squared distance for each point as x² + y².
   • Push the point and its distance into the heap.
   • If the heap size exceeds k, remove the farthest point.
3. Once all points are processed, the heap contains the k closest points.
4. Extract and return these k points.

---

Java Code Implementation

import java.util.*;

public class KClosestPointsToOrigin {

    public int[][] kClosest(int[][] points, int k) {
        // Step 1: Create a max-heap based on squared distance
        PriorityQueue<int[]> maxHeap = new PriorityQueue<>(
            (a, b) -> Integer.compare(distance(b), distance(a))
        );

        // Step 2: Process each point
        for (int[] point : points) {
            maxHeap.offer(point); // Add point to heap
            if (maxHeap.size() > k) {
                maxHeap.poll(); // Remove the farthest point
            }
        }

        // Step 3: Extract k closest points
        int[][] result = new int[k][2];
        for (int i = 0; i < k; i++) {
            result[i] = maxHeap.poll();
        }

        return result;
    }

    // Helper method to calculate squared distance from origin
    private int distance(int[] point) {
        return point[0] * point[0] + point[1] * point[1];
    }

    // Optional main method for testing
    public static void main(String[] args) {
        KClosestPointsToOrigin obj = new KClosestPointsToOrigin();

        int[][] points1 = {{1,3}, {-2,2}};
        int k1 = 1;
        System.out.println(Arrays.deepToString(obj.kClosest(points1, k1))); // [[-2, 2]]

        int[][] points2 = {{3,3}, {5,-1}, {-2,4}};
        int k2 = 2;
        System.out.println(Arrays.deepToString(obj.kClosest(points2, k2))); // [[3, 3], [-2, 4]]
    }
}

---

Dry Run Example

Input:
points = [[3,3], [5,-1], [-2,4]], k = 2

Step-by-step:

1. Start with an empty heap.
2. Add [3,3] → distance = 18. Heap = [[3,3]]
3. Add [5,-1] → distance = 26. Heap = [[5,-1], [3,3]]
4. Add [-2,4] → distance = 20. Heap = [[5,-1], [3,3], [-2,4]]
   Heap size exceeds k = 2, remove the farthest → [5,-1] (distance 26).
   Remaining heap = [[3,3], [-2,4]]
5. Extract heap contents → [[3,3], [-2,4]]

Output: [[3,3], [-2,4]]

---

Complexity Analysis

• Time Complexity:
  O(n log k)
  • Each point insertion/removal in the heap takes O(log k) time.
  • For n points, total operations = O(n log k).
• Space Complexity:
  O(k)
  The heap stores at most k points at any time.

If you use the sorting-based approach, the time complexity would be O(n log n) and space complexity O(n).

---

Alternative Approaches

1. Sorting-Based Approach:
   • Compute all distances and sort all points.
   • Return the first k points.
   • Simpler to implement but slower for large n.
     Time complexity: O(n log n).
2. QuickSelect Algorithm (Average O(n)):
   • Use QuickSelect (similar to quicksort partitioning) to find the kth smallest distance.
   • Partition points so that all smaller distances come before the kth distance.
   • Faster average case but complex to implement and less stable in worst-case scenarios.