Print X array elements closest to the Kth smallest element in the array

Given two integers K, X, and an array arr[] consisting of N distinct elements, the task is to find X elements closest to the K^th smallest element from the given array.

Examples:

Input: arr[] = {1, 2, 3, 4, 10}, K = 3, X = 2
Output: 2 3
Explanation: K^th smallest element present in the given array is 3 and X(= 2) closest array elements to 3 are {2, 3} or {3, 4}.
Therefore, the required output is 2 3.

Input: arr[] = {1, 9, 8, 2, 7, 3, 6, 4, 5, 10, 13, 12, 16, 14, 11, 15}, K = 3, X = 5
Output: 1 2 3 4 5

Naive Approach: The simplest approach to solve this problem is to sort the array and print X closest elements to the K^th smallest element of the given array using the two-pointers technique.
Time Complexity: O(N * log N)
Auxiliary Space: O(1)

Efficient Approach: To optimize the above approach the idea is to efficiently compute the value of the K^th smallest element of the given array using the Median Selection Algorithm.  Follow the steps below to solve the problem:

• Calculate the Kth smallest, say KthElem of the given array using the Median Selection Algorithm.
• Initialize an array, say diff[] to store the absolute difference of arr[i] and KthElem.
• Create a map, say maps to map each element of the array to the absolute difference of the current element and KthElem.
• Traverse the given array and append arr[i] to maps[abs(KthElem – arr[i])].
• Calculate X^th smallest element, say XthElem of diff[] array using the Median Selection Algorithm to print exactly X closest items.
• Finally, traverse the diff[] array and check if XthElem less than or equal to diff[i] or not. If found to be true then print the element of the array with the help of maps.

Below is the implementation of the above approach:

C++

// C++ program to implement
// the above approach
#include <bits/stdc++.h>
#include <vector>
#include <unordered_map>
#include <algorithm>

using namespace std;
// Function to swap
// two elements of array
void swap(vector<int>& arr, int a, int b) {
    int temp = arr[a];
    arr[a] = arr[b];
    arr[b] = temp;
}

// Function to partition
// the array around x
int partition(vector<int>& arr, int l, int r, int x) {

    // Traverse array
    // from index l to r
    for (int i = l; i < r; i++) {

        // partition array
        // around x
        if (arr[i] == x) {
            swap(arr, r, i);
            break;
        }
    }
    x = arr[r];
    int i = l;

    // Traverse array
    // from index l to r
    for (int j = l; j < r; j++) {
        if (arr[j] <= x) {
            swap(arr, i, j);
            i++;
        }
    }
    swap(arr, i, r);
    return i;
}

// Function to find
// median of arr[]
// from index l to l + n
int findMedian(vector<int>& arr, int l, int n) {
    vector<int> lis;
    for (int i = l; i < l + n; i++) {
        lis.push_back(arr[i]);
    }

    // Sort the array
    sort(lis.begin(), lis.end());

    // Return middle element
    return lis[n / 2];
}

// Function to get
// the kth smallest element
int kthSmallest(vector<int>& arr, int l, int r, int k) {

    // If k is smaller than
    // number of elements
    // in array
    if (k > 0 && k <= r - l + 1) {

        // Stores count of
        // elements in arr[l..r]
        int n = r - l + 1;


        // Divide arr[] in groups
        // of size 5, calculate
        // median  of every group
        // and store it in
        // median[] array.
        vector<int> median;
        int i = 0;
        while (i < n / 5) {
            median.push_back(findMedian(arr, l + i * 5, 5));
            i++;
        }

        // For last group with
        // less than 5 elements
        if (i * 5 < n) {
            median.push_back(findMedian(arr, l + i * 5, n % 5));
            i++;
        }
        int medOfMed;

        // If median[] has
        // only one element
        if (i == 1) {
            medOfMed = median[i - 1];
        }

        // Find median of all medians
        // using recursive call.
        else {
            medOfMed = kthSmallest(median, 0, i - 1, i / 2);
        }

        // Stores position of pivot
        // element in sorted array
        int pos = partition(arr, l, r, medOfMed);

        // If position is same as k
        if (pos - l == k - 1) {
            return arr[pos];
        }

        // If position is more,
        if (pos - l > k - 1) {

            // recur for left subarray
            return kthSmallest(arr, l, pos - 1, k);
        }
        // Else recur for right subarray
        return kthSmallest(arr, pos + 1, r, k - pos + l - 1);
    }
    // If k is more than
    // number of elements
    // in the array
    return INT_MAX;
}

// Function to print
vector<int> closestElements(vector<int>& arr, int k, int x) {

    // Stores size of arr
    int n = arr.size();

    // Stores kth smallest
    // of the given array
    int KthElem = kthSmallest(arr, 0, n - 1, k);

    // Store the value of
    // abs(KthElem - arr[i])
    vector<int> diff;

    // Create a map to map
    // array element to
    // abs(KthElem - arr[i])
    unordered_map<int, vector<int>> maps;
    for (int elem : arr) {

        // Stress the value of
        // abs(elem - KthElem)
        int temp = abs(elem - KthElem);

        // map array elements
        maps[temp].push_back(elem);

        // append temp
        diff.push_back(temp);
    }
    int XthElem = kthSmallest(diff, 0, n - 1, x);
    vector<int> res;
    for (int dx : diff) {
        if (dx <= XthElem) {
            for (int elem : maps[dx]) {
                if (res.size() < x) {
                    res.push_back(elem);
                }
            }
        }
    }
    return res;
}
// Driver Code
int main()
{
    vector<int> arr = {1, 2, 3, 4, 10, 15};
    int k = 3, x = 2;

    // Store X closest elements
    auto res = closestElements(arr, k, x);

    // Print X closest elements
    for (auto i : res) {
        cout << i << " ";
    }
    return 0;
}

// This code is contributed by shivhack999

Python3

# Python3 program to implement
# the above approach

import collections


# Function to swap
# two elements of array
def swap(arr, a, b):
    temp = arr[a]
    arr[a] = arr[b]
    arr[b] = temp

# Function to partition
# the array around x
def partition(arr, l, r, x):

    # Traverse array
    # from index l to r
    for i in range(l, r):

        # partition array
        # around x
        if arr[i] == x:
            swap(arr, r, i)
            break

    x = arr[r]
    i = l

    # Traverse array
    # from index l to r
    for j in range(l, r):
        if (arr[j] <= x):
            swap(arr, i, j)
            i += 1
    swap(arr, i, r)
    return i

# Function to find
# median of arr[]
# from index l to l + n
def findMedian(arr, l, n):
    lis = []
    for i in range(l, l + n):
        lis.append(arr[i])

    # Sort the array
    lis.sort()

    # Return middle element
    return lis[n // 2]

# Function to get
# the kth smallest element
def kthSmallest(arr, l, r, k):

    # If k is smaller than
    # number of elements
    # in array
    if (k > 0 and
        k <= r - l + 1):

        # Stores count of
        # elements in arr[l..r]
        n = r - l + 1

        # Divide arr[] in groups
        # of size 5, calculate
        # median  of every group
        # and store it in
        # median[] array.
        median = []

        i = 0
        while (i < n // 5):
            median.append(
            findMedian(arr,
                l + i * 5, 5))
            i += 1

        # For last group with
        # less than 5 elements
        if (i * 5 < n):
            median.append(
             findMedian(arr,
                l + i * 5, n % 5))
            i += 1


        # If median[] has
        # only one element
        if i == 1:
            medOfMed = median[i - 1]

        # Find median of all medians
        # using recursive call.
        else:
            medOfMed = kthSmallest(
             median, 0, i - 1, i // 2)

        # Stores position of pivot
        # element in sorted array
        pos = partition(arr, l, r,
                         medOfMed)

        # If position is same as k
        if (pos - l == k - 1):
            return arr[pos]

        # If position is more,
        if (pos - l > k - 1):

            # recur for left subarray
            return kthSmallest(arr, l,
                          pos - 1, k)

        # Else recur for right subarray
        return kthSmallest(arr, pos + 1,
                    r, k - pos + l - 1)

    # If k is more than
    # number of elements
    # in the array
    return 999999999999

# Function to print
def closestElements(arr, k, x):

    # Stores size of arr
    n = len(arr)

    # Stores kth smallest
    # of the given array
    KthElem = kthSmallest(
            arr, 0, n - 1, k)

    # Store the value of
    # abs(KthElem - arr[i])
    diff = []

    # Create a map to map
    # array element to
    # abs(KthElem - arr[i])
    maps = collections.defaultdict(
                              list)
    for elem in arr:

        # Stress the value of
        # abs(elem - KthElem)
        temp = abs(elem - KthElem)

        # map array elements
        maps[temp].append(elem)

        # append temp
        diff.append(temp)

    XthElem = kthSmallest(diff, 0,
                        n - 1, x)

    # Store X closest elements
    res = set()

    # Traverse diff[] array
    for dx in diff:

        # If absolute difference is
        # less than or equal to XthElem
        if dx <= XthElem:

            # Append closest elements
            for elem in maps[dx]:
                if len(res) < x:
                  res.add(elem)
    return res


# Driver Code
if __name__ == '__main__':

    arr = [1, 2, 3, 4, 10, 15]
    k = 3
    x = 2

    # Store X closest elements
    res = closestElements(arr, k, x)

    # Print X closest elements
    for i in res:
        print(i, end =" ");

Javascript

// JavaScript program to implement
// the above approach

// Function to swap
// two elements of array
function swap(arr, a, b) {
  let temp = arr[a];
  arr[a] = arr[b];
  arr[b] = temp;
}

// Function to partition
// the array around x
function partition(arr, l, r, x) {

// Traverse array
// from index l to r
  for (let i = l; i < r; i++) {

      // partition array
    // around x
    if (arr[i] == x) {
      swap(arr, r, i);
      break;
    }
  }

  x = arr[r];
  let i = l;

  // Traverse array
  // from index l to r
  for (let j = l; j < r; j++) {
    if (arr[j] <= x) {
      swap(arr, i, j);
      i += 1;
    }
  }

  swap(arr, i, r);
  return i;
}

// Function to find
// median of arr[]
// from index l to l + n
function findMedian(arr, l, n) {
  let lis = [];
  for (let i = l; i < l + n; i++) {
    lis.push(arr[i]);
  }

  // Sort the array
  lis.sort();

  // Return middle element
  return lis[(Math.floor(n / 2))];
}

// Function to get
// the kth smallest element
function kthSmallest(arr, l, r, k) {
  // If k is smaller than
  // number of elements
  // in array
  if (k > 0 && k <= r - l + 1) {

      // Stores count of
    // elements in arr[l..r]
    let n = r - l + 1;

    // Divide arr[] in groups
    // of size 5, calculate
    // median of every group
    // and store it in
    // median[] array.
    let median = [];
    let i = 0;

    while (i < Math.floor(n / 5)) {
      median.push(findMedian(arr, l + i * 5, 5));
      i += 1;
    }


    // For last group with
    // less than 5 elements
    if (i * 5 < n) {
      median.push(findMedian(arr, l + i * 5, n % 5));
      i += 1;
    }

    let medOfMed;

    // If median[] has
    // only one element
    if (i == 1) {
      medOfMed = median[i - 1];
    } else {
        // Find median of all medians
        // using recursive call.
      medOfMed = kthSmallest(median, 0, i - 1, Math.floor(i / 2));
    }

    // Stores position of pivot
    // element in sorted array
    let pos = partition(arr, l, r, medOfMed);

    // If position is same as k
    if (pos - l == k - 1) {
      return arr[pos];
    }

    // If position is more,
    if (pos - l > k - 1) {
      return kthSmallest(arr, l, pos - 1, k);
    }

    return kthSmallest(arr, pos + 1, r, k - pos + l - 1);
  }

      // If k is more than
    // number of elements
    // in the array
  return 999999999999;
}

function closestElements(arr, k, x) {
  let n = arr.length;

    // Stores kth smallest
    // of the given array
  let KthElem = kthSmallest(arr, 0, n - 1, k);

  // Store the value of
  // abs(KthElem - arr[i])
  let diff = [];
  let maps = {};

  for (let elem of arr) {
      // Stress the value of
    // abs(elem - KthElem)
    let temp = Math.abs(elem - KthElem);

    if (maps[temp] === undefined) {
      maps[temp] = [];
    }

    maps[temp].push(elem);
    diff.push(temp);
  }

  let XthElem = kthSmallest(diff, 0, n - 1, x);
  let res = new Set();

  for (let dx of diff) {
    if (dx <= XthElem) {
      for (let elem of maps[dx]) {
        if (res.size < x) {
          res.add(elem);
        }
      }
    }
  }
  const finalAns = [...res];
  finalAns.sort();

  return finalAns;
}

// Driver Code
let arr = [1, 2, 3, 4, 10, 15];
let k = 3;
let x = 2;
let res = closestElements(arr, k, x);

// Print X closest elements
console.log(res.join(" "));

Output:

2 3

Time Complexity: O(N)
Auxiliary Space: O(N)