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 Kth smallest element from the given array.

Examples:

Input: arr[] = {1, 2, 3, 4, 10}, K = 3, X = 2
Output: 2 3
Explanation: Kth 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 Kth 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 Kth smallest element of the given array using the Median Selection Algorithm.  Follow the steps below to solve the problem:



Below is the implementation of the above approach:

filter_none

edit
close

play_arrow

link
brightness_4
code

# 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:
          
        # Stres 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 eqaul 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 =" ");
     
chevron_right

Output:
2 3

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

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready.





Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.


Article Tags :