Find the K Closest Points to Origin using Priority Queue

Introduction

Priority Queue is an important data structure. It is implemented using Heap Data Structure and has a wide range of applications. It is also important when it comes to technical interviews as well. This blog will discuss the solution to this problem to find the k closest points to the origin using a priority queue. So before we deep dive into the answer, we should look at some examples to understand the problem better.

Some Examples

Input:

 [(1, 0), (2, 2), (1,3), (0,-4)]   K=2

Output: 

[(1, 0), (2, 2)]

Explanation:

• The Square of Distance of (1, 0) from the origin is 1
• The Square of Distance of (1, 3) from the origin is 10
• The Square of Distance of (2, 2) from the origin is 8
• The Square of Distance of (0, -4) from the origin is 16 

Therefore for K = 2, the closest points are (1, 0) & (2, 2)

Input:

[(2, 1), (-3, -2), (3, 0), (0, -2), (3, -1)]   K = 3

Output:

[(0, -2), (2, 1), (3, 0)] 

Explanation:

• The Square of Distance of (2, 1) from the origin is 5
• The Square of Distance of (-3, -2) from the origin is 13
• The Square of Distance of (3, 0) from the origin is 9
• The Square of Distance of (0, -2) from the origin is 4
• The Square of Distance of (3, -1) from the origin is 10

Therefore for K = 3, the closest points are (0, -2), (2, 1) & (3, 0)

*Square of Distance = x*x + y*y

Also Read, Prefix Sum Array

Priority Queue Approach

To solve this problem, find the K closest points to the origin using the priority queue; we will first create a min-heap of pairs in which we will store the distance of the point from the origin and the point itself, and after that, we will traverse the min-heap till K. Simultaneously we will store all the points in our final array. After that, we will return the array.

Algorithm

1. To solve this problem, find the K closest points to the origin using the priority queue. We first need to create a function kClosestPoints() that takes three parameters: the array, size, and value K. 
2. In the function, we will declare a priority queue of pairs that will store the distance and the coordinates of points in increasing order. 
3. We will now calculate the distance of every point from the origin, and then we will store the distance and the coordinates of a point in the priority queue.
4.  After that, we will pick the priority queue’s first K elements, store them in the final array, and then return them.

Code in C++

// C++ code for Find the K closest points to origin using Priority Queue
#include <bits/stdc++.h>
using namespace std;

// function to find the K closest points from the origin
vector<pair<int, int>> kClosestPoints(vector<pair<int, int>> points, int n, int k)
{
    // to store the answer
    vector<pair<int, int>> ans;

   // to store the in an increasing order
    // according to their distance from the origin
    priority_queue<int, vector<pair<int, pair<int, int>>>, greater<pair<int, pair<int, int>>>> p;

    for (int i = 0; i < n; i++)
    {
        // calculating the distance of the point from the origin
        int distance = points[i].first * points[i].first + points[i].second * points[i].second;
        // storing the distance with the point in the priority queue
        p.push({distance, {points[i].first, points[i].second}});
    }

   // taking out the K closest points from the
    // priority queue and storing them
    while (k--)
    {
        auto x = p.top();
        ans.push_back({x.second.first, x.second.second});
        p.pop();
    }
    return ans;
}

// Driver Code
int main()
{

    int k = 3, n = 6;

    vector<pair<int, int>> arr{{-9, -2}, {-10, 1}, {3, -8}, {4, 4}, {-9, 0}, {7, 7}};

    vector<pair<int, int>> ans = kClosestPoints(arr, n, k);
  
    for (int i = 0; i < k; i++)
    {
        cout << "{" << ans[i].first << "," << ans[i].second << "} ";
    }
}

You can also try this code with Online C++ Compiler

Run Code

Input and Output:

Input: N = 6 , K = 3
[(-9, -2), (-10, 1), (3, 8), (4, 4), (-9, 0), (7, 7)]

Output: {4, 4}, {3, -8}, {-9, 0}

Complexity Analysis

Time Complexity: O(N)

Since we are doing single traversals, therefore, the time complexity is O(N).

Space Complexity: O(N)

Since we store our answer in another array, the space complexity will be O(N).