Heap: k closest number

Problem Statement:

You are given a sorted array and a number k and a number x.

You need to return k elements that are closest to x

Example:

Input: arr = [1, 2, 3, 4, 5] k = 3 x = 3

Output:  [1, 2, 3]

Solution Explanation:

We will use max heap to solve the problem.

We will store a pair [absolute difference with x, corresponding key]

When the size of the heap is more than k, then pop it.

Return the elements from the heap

Time Complexity: O(logn)
Space Complexity: O(n)

Code Solution

#include <algorithm>
#include <iostream>
#include <vector>
#include <queue>
using namespace std;

vector<int> solution(vector<int>& arr, int k, int x) 
{
    
    priority_queue<pair<int,int>> pq;
    
    for(int i=0;i<arr.size();i++)
    {
    	pq.push({abs(arr[i]-x),arr[i]});

      	if(pq.size()>k)
          pq.pop();
    }

    vector<int> ans;

    while(pq.size()>0)
    {
        ans.push_back(pq.top().second);
        pq.pop();
    }

    sort(ans.begin(),ans.end());
    return ans;
}

int main() 
{
    vector<int> arr = {1, 2, 3, 4, 5};
    int k = 3;
    int x = 3;

    vector<int> sol = solution(arr, k, x);
    
    for (int num : sol) 
    {
        cout << num << " ";
    }

    return 0;
}

Output

2 3 4