K Closest Points to Origin (Medium)

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).

Example 1:

Input: points = [[1,3],[-2,2]], k = 1
Output: [[-2,2]]
Explanation:
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]]
Explanation: The answer [[-2,4],[3,3]] would also be accepted.

1) Explain the problem

You are given an array of points on an X-Y plane and an integer k. Each point is represented as [xi, yi]. Your task is to find the k closest points to the origin (0, 0) based on their Euclidean distance.

2) Short easy to remember solution/approach

Use a max heap (priority queue) to keep track of the k closest points to the origin. Calculate the Euclidean distance for each point and maintain the k smallest distances in the heap.

3) Solution in JavaScript with code commenting

// Helper function to calculate the Euclidean distance from the origin
function euclideanDistance(point) {
    return Math.sqrt(point[0] * point[0] + point[1] * point[1]);
}

// Helper function to create a max heap
class MaxHeap {
    constructor() {
        this.heap = [];
    }

    // Insert a new element into the heap
    insert(val) {
        this.heap.push(val);
        this._heapifyUp();
    }

    // Remove and return the max element (root) of the heap
    extractMax() {
        if (this.size() === 0) return null;
        if (this.size() === 1) return this.heap.pop();
        const max = this.heap[0];
        this.heap[0] = this.heap.pop();
        this._heapifyDown();
        return max;
    }

    // Get the size of the heap
    size() {
        return this.heap.length;
    }

    // Get the max element (root) of the heap without removing it
    peek() {
        return this.heap[0];
    }

    // Internal method to maintain heap property after insertion
    _heapifyUp() {
        let index = this.heap.length - 1;
        while (index > 0) {
            let parentIndex = Math.floor((index - 1) / 2);
            if (this.heap[index][0] <= this.heap[parentIndex][0]) break;
            [this.heap[index], this.heap[parentIndex]] = [this.heap[parentIndex], this.heap[index]];
            index = parentIndex;
        }
    }

    // Internal method to maintain heap property after extraction
    _heapifyDown() {
        let index = 0;
        while (true) {
            let leftChildIndex = 2 * index + 1;
            let rightChildIndex = 2 * index + 2;
            let largest = index;

            if (leftChildIndex < this.size() && this.heap[leftChildIndex][0] > this.heap[largest][0]) {
                largest = leftChildIndex;
            }
            if (rightChildIndex < this.size() && this.heap[rightChildIndex][0] > this.heap[largest][0]) {
                largest = rightChildIndex;
            }
            if (largest === index) break;
            [this.heap[index], this.heap[largest]] = [this.heap[largest], this.heap[index]];
            index = largest;
        }
    }
}

// Function to find the k closest points to the origin
function kClosest(points, k) {
    const maxHeap = new MaxHeap();

    // Insert the first k points into the max heap
    for (let i = 0; i < k; i++) {
        maxHeap.insert([euclideanDistance(points[i]), points[i]]);
    }

    // Iterate over the remaining points
    for (let i = k; i < points.length; i++) {
        const dist = euclideanDistance(points[i]);
        if (dist < maxHeap.peek()[0]) {
            maxHeap.extractMax();
            maxHeap.insert([dist, points[i]]);
        }
    }

    // Extract the k closest points from the max heap
    const result = [];
    while (maxHeap.size() > 0) {
        result.push(maxHeap.extractMax()[1]);
    }

    return result;
}

4) Explanation of the solution in an easy-to-understand way

Imagine you have a bunch of points on a map, and you want to find the ones closest to a specific spot (the origin). You can use a special basket that always keeps the closest points on top. You keep adding points to this basket and always check if the new point is closer than the farthest point in the basket. In the end, you take out the points from the basket, and these are your closest points.

5) Code explanation in pointers

euclideanDistance function calculates the distance of a point from the origin using the Euclidean distance formula.
MaxHeap class creates a max heap to manage the k closest points efficiently.
kClosest function:
- Inserts the first k points into the max heap.
- Iterates over the remaining points, replacing the farthest point in the heap if the current point is closer.
- Extracts the k closest points from the max heap and returns them.

6) Complexities

Time Complexity: O(n log k), where n is the number of points and k is the number of closest points to find. This is because inserting and extracting from the heap takes O(log k) time, and we do this for each of the n points.
Space Complexity: O(k), where k is the number of closest points to find. This is due to the space required to store the k closest points in the heap.