Find the kth Smallest Element by Partitioning the Array

This article finds the kth smallest element in an array by partitioning it.

Consider the following example.

Array= 33, 1, 0, 15, 16, 12, 14, 11, 9, 8, 2
3rd smallest element = 2
5th smallest element = 9

Naïve Approach

One approach to solving that problem would be to sort the array first, let’s say, using the quick sort algorithm, and then access the element at index (k-1). Here, we assume that the array starts from index 0. This approach might be intuitive and straightforward. However, its average time complexity is O(nlogn), where n is the array size. We can find solutions that have better complexities.

Quick Select Algorithm

As we only want the kth smallest element, there is no need to sort the entire array. The quick select algorithm takes advantage of this idea. Let’s see how it works.

  • First, it finds a random pivot element from the array.
  • Then, it partitions the array based on that pivot point. Values that are less than the pivot element get shifted to the left, and items greater than that get shifted to the right.
  • If the index of the pivot is less than (k-1), then it means our required element lies on the right side of the array. So, only that part gets considered for the next call of the algorithm, i.e., the next pivot will be from this side.
  • If the pivot’s position is more than (k-1), then only the left part gets considered.
  • If our required (kth smallest) element is equal to the pivot element, then the recursion stops, and that value gets returned.

The algorithm runs until the pivot position is equal to (k-1).

As you can observe, this algorithm works like quicksort. However, it recurs for one side only. Therefore, its average time complexity is better than the quick sort algorithm, i.e., O(n). The worst-case complexity is still O(n2).

Implementation

Consider the following code to find the kth smallest element.

import math
import random

def partition(array, l, r, index):
  pivot = array[index] #find the pivot element
  left=[] #elements less than or equal to pivot
  right=[] #elements greater than pivot
  for i in range(l, r+1): 
    if i==index: #leave the pivot element
      pass
    elif array[i] <= pivot:
      left.append(array[i])
    else:
      right.append(array[i])

  #make changes in the actual array
  end_left = l + len(left)
  array[l: end_left] = left 
  array[end_left] = pivot
  array[end_left+1:r+1] = right

  return end_left

def selectPivot(array, l, r): #find a random pivot index
  return random.randint(l, r)

def quickSelect(array, l, r, k):
  if l <= r:
    index = selectPivot(array, l, r) #get the index of the pivot point
    pivot_pos = partition(array,l, r,  index) #partition the array based on the pivot
    if pivot_pos==k: #if pivot is the required smallest element then return it
      return array[k]
    elif pivot_pos < k: #consider the right part
        return quickSelect(array,pivot_pos+1, r, k)
    else:
        return quickSelect(array, l, pivot_pos-1, k) #consider the left part

array=[34, 1, 4, 2, 18, 29, 14, 28, 23]
k=3
kth_smallest = quickSelect(array, 0, len(array)-1, k-1)
print(f"{k} smallest element is: {kth_smallest}")
k=4
kth_smallest = quickSelect(array, 0, len(array)-1, k-1)
print(f"{k} smallest element is: {kth_smallest}")

Output

3 smallest element is: 4
4 smallest element is: 14

Since we are choosing a random point, we can get unlucky as well. Let’s explain this using an example. We have an array A = [0, 5, 3, 1, 4], and we want to find the 1st smallest element. At first, we get 5 as a pivot. The partitioned array is now [0, 3, 1, 4] as all the elements fall to the left. Next time, we get 4, and then the partitioned array will be [0, 3, 1]. After that, we get 3, 1, and 0 as pivot elements. As you can observe, the algorithm runs 5+4+3+2+1 times or n+n-1+…+1 times, which becomes O(n2).

Median of Medians Algorithm

Selecting a pivot in such a way that never leads to the above-given situation can guarantee a time complexity of O(n) in the worst-case scenario as well. For this, we have the median of medians algorithm.
This algorithm works as follows:

  • It divides an array into chunks, which are usually of size 5. The last element can have less than 5 items.
  • Then, it finds medians of these chunks by sorting and store results in an array.
  • If only one median is left, then it gets returned as a pivot element. Otherwise, the algorithm works again on the medians array until we get a single median.

Code

Let’s find the kth smallest element using the median of medians method.

import math
import random

def partition(array, l, r, pivot):
  index = array.index(pivot)
  left=[] #elements less than or equal to pivot
  right=[] #elements greater than pivot
  for i in range(l, r+1): 
    if i==index: #leave the pivot element
      pass
    elif array[i] <= pivot:
      left.append(array[i])
    else:
      right.append(array[i])

  #make changes in the actual array
  end_left = l + len(left)
  array[l: end_left] = left 
  array[end_left] = pivot
  array[end_left+1:r+1] = right

  return end_left

def medianOfMedians(arr, l, r):
  groups=[]
  #create chunks of size 5
  temp = arr[l:r+1]
  for i in range(0, len(temp), 5):
      groups.append(temp[i:i+5])
  medians =[]

  #find median of every chunk
  for group in groups:
    medians.append(sorted(group)[(len(group))//2])
  l_med = len(medians)
  if l_med == 1:
    return medians[0]
  else:
    return medianOfMedians(medians, 0, l_med-1)

def quickSelect(array, l, r, k):
  if l <= r:
    pivot = medianOfMedians(array, l, r) #get the index of the pivot point
    pivot_pos = partition(array,l, r,  pivot) #partition the array based on the pivot
    if pivot_pos==k: #if pivot is the required smallest element then return it
      return array[k]
    elif pivot_pos < k: #consider the right part
      return quickSelect(array,pivot_pos+1, r, k)
    else:
      return quickSelect(array, l, pivot_pos-1, k) #consider the left part


arr=[34, 1, 4, 2, 18, 29, 14, 28, 23]
k=3
kth_smallest = quickSelect(arr, 0, len(arr)-1, k-1)
print(f"{k} smallest element is: {kth_smallest}")
k=4
kth_smallest = quickSelect(arr, 0, len(arr)-1, k-1)
print(f"{k} smallest element is: {kth_smallest}")

Output

3 smallest element is: 4
4 smallest element is: 14

Back to: Data Structure > Data Structure - Arrays

Leave a Reply

Your email address will not be published. Required fields are marked *