AlgoMaster Sheet – Find K Closest Elements

Problem Statement: Given a sorted integer array arr, two integers k and x, return the k closest integers to x in the array. The result should also be sorted in ascending order.

An integer a is closer to x than an integer b if:

  • |a - x| < |b - x|, or
  • |a - x| == |b - x| and a < b

Example 1:

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

Example 2:

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

Solution:

Disclaimer: Don’t jump directly to the solution, try out yourself first.

Practice Questions   Find K Closest Elements   
Company Tags: Amazon

Solution 1: 

Algorithm:

  • Initially let ‘low’ = 0 and ‘high’ = N – 1, Where ‘low’ and ‘high’ tell the range of integers taken in the final answer.
  • If at any point (high – low + 1) = K, i.e., the size of the range is ‘K’, that range of integers is the answer. And we create a new array of that range and return it.
  • If not, then we check if which one of ‘A[low]’ and ‘A[high]’ is at a smaller distance from ‘X’, whoever is at a smaller distance we keep it, and remove the other.
  • If ‘low’ is at a smaller distance, we decrement ‘high’, i.e., high = high – 1.
  • If ‘high’ is at a smaller distance, we increment ‘low’, i.e., low = low + 1.

Time Complexity: O(N), where N is the size of the array.
Space Complexity: O(N), since no additional space is used.

 Where ‘N’ is length of the input array.

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

vector<int> findClosestElements(vector<int>& arr, int k, int x) {
    int low = 0, high = arr.size()- 1;

    while (high - low + 1 > k) {
        if (abs(arr[low] - x) <= abs(arr[high] - x)) {
            high--;
        } else {
            low++;
        }
    }

    vector<int> ans;
    for (int i = low; i <= high; i++) {
        ans.push_back(arr[i]);
    }

    return ans;
}

int main()
{
    vector<int> input = {10, 12, 15, 17, 18, 20, 25};
    int target = 16, k = 4;

    vector<int> result = findClosestElements(input, k, target);
    for (int i : result)
    {
        cout << i << " ";
    }

    return 0;
}


Output:

12 15 17 18



Solution 2: 

We can also find the insertion point i using the Binary search Algorithms, which runs in O(log(n)) time. Since finding the k closest elements takes O(k) time, the overall time complexity of this solution is O(log(n) + k).

Time Complexity: O(log(n) + k).

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

// Function to find the `k` closest elements to `target` in a sorted integer vector `input`
vector<int> findClosestElements(vector<int>& arr, int k, int x) {

    // find the insertion point using the binary search algorithm
    int i = lower_bound(arr.begin(), arr.end(), x) - arr.begin();
    int left = i - 1;
    int right = i;

    // run `k` times
    while (k > 0)
    {
        // compare the elements on both sides of the insertion point `i`
        // to get the first `k` closest elements

        if (left < 0 || (right < arr.size() &&
                         abs(arr[left] - x) > abs(arr[right] - x)))
        {
            right++;
        }
        else
        {
            left--;
        }
        k--;
    }

    // return `k` closest elements
    return vector<int>(arr.begin() + left + 1, arr.begin() + right);
}

int main()
{
    vector<int> input = {10, 12, 15, 17, 18, 20, 25};
    int target = 16, k = 4;

    vector<int> result = findClosestElements(input, k, target);
    for (int i : result)
    {
        cout << i << " ";
    }

    return 0;
}

Output:

12 15 17 18

This article is contributed by : Sayan Dutta to this article on csforall.in

Leave a Comment