Table of Contents
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