Heap: K Closest Points To Origin

Problem Statement:

You are given a 2D array and a value k.

You need to return the k closest to the origin (0,0).

Example:

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

Solution 1: Sorting Approach

The mathematical formula to get the point closer to the origin is using euclidean distance.

For euclidean distance, we have to perform the square root, but it will make the code more complex.

Hence we can simply use the following formula without using the square root `dist = x*x+y*y`

Time Complexity: O(nlogK)
Space Complexity: O(1)

Solution 2: Max Heap

We will use max heap to solve the problem.

Take priority queue to store in maxHeap.

Calculate the new points and maintain the max heap of size k.

Then if the size is greater than k, then remove the root element. Hence the max heap is always have the top k smallest elements.

Time Complexity: O(NlogK)
Space Complexity: O(1)

Code Solution

#include <iostream>
#include <vector>
#include <algorithm>
#include <queue>

using namespace std;

vector<vector<int>> solution_1(vector<vector<int>>& arr, int k) 
{
    sort(arr.begin(), arr.end(), [](vector<int>& a, vector<int>& b) {
        return a[0] * a[0] + a[1] * a[1] < b[0] * b[0] + b[1] * b[1];
    });
    return vector<vector<int>>(arr.begin(), arr.begin() + k);
}

vector<vector<int>> solution_2(vector<vector<int>>& arr, int k) 
{

    vector<vector<int>> result(k);
    priority_queue<vector<int>> maxHeap;
    
    for (auto& p : arr) 
    {
        int x = p[0], y = p[1];
        maxHeap.push({x*x + y*y, x, y});

        if (maxHeap.size() > k) 
        {
            maxHeap.pop();
        }
    }
    
    for (int i = 0; i < k; ++i) 
    {
        vector<int> top = maxHeap.top();
        maxHeap.pop();
        result[i] = {top[1], top[2]};
    }
    
    return result;
}

int main() 
{
    vector<vector<int>> points ={{1, 3}, {-2, 2}};
    int k = 1;
    
    vector<vector<int>> res = solution_1(points, k);
    
    cout << "Solution 1:"<<endl;

     for (vector<int> point : res) 
     {
        cout << point[0] << ", " << point[1];
        cout<<endl;
    }

    res = solution_2(points, k);

    cout << "Solution 2:"<<endl;

     for (vector<int> point : res) 
     {
        cout << point[0] << ", " << point[1];
        cout<<endl;
    }
    return 0;
}

Output

Solution 1:
-2, 2
Solution 2:
-2, 2