Implementing Quick Select Algorithm - Finding kth largest/smallest element
Today I am going to write about a special type of algorithm for finding a kth largest or smallest element in the array. This is my favorite data structure question due to intricacy involved in solving the problem. There are many variants of this problem including more efficient using max/min heap. But I am going to pass that solution and prefer quick-select algorithm solely for the purpose of simplicity.
For the sake of this post, I am going to use Swift as a primary programming language for this exercise, but algorithm can be implemented in any programming language
Quick select is a special form of quick sort algorithm. However, instead of sorting the full array, it partially sorts it until it finds a desired pivot position. For example, if you input an array to quick-select and ask it to find kth smallest/largest element, as soon as it finds the pivot index equal to k, it returns an element at that index as an answer.
Algorithm can be summarized as follows,
For the sake of simplicity, I am going to use the case to find kth smallest element. With slight change we can use it to find kth largest element with similar flow
Source: https://carbon.now.sh/
Please note that indices low and high are inclusive of both ends. Which means for the initial iteration low = 0 and high = arraySize - 1
To help better understand the code, let's dry run it on random input and see how it outputs. For the sake of example, our input would be
var input = [100, 300, 2, 3, 1000, 67, 89, 10, 764, 1, 546]
At the start, input will be our input array, with low = 0 and high = 10 and value of k = 3 that is we want to find kth smallest element.
Please note that this is an in-place algorithm. Which means as it finds the element, it also reorders element. So array used on first iteration may not be exactly same as the one used during next iteration
First iteration,
k = 3
-- – – – – – – – – -- – – – – – – – – -- – – – – – – – – -- -- -- - --
| 100 | 300 | 2 | 3 | 67 | 89 | 10 | 1 | 546 | 1000 | 764 |
-- – – – – – – – – -- – – – – – – – – -- – – – – – – – – -- -- -- - --
⬆ ⬆ ⬆
low pivot high
low != high and pivot - low != k hence continue to next iteration
Second iteration
k = 3
-- – – – – – – – – -- – – – – – – – – -- – – – – – – – – -- -- -- - -
| 1 | 300 | 2 | 3 | 67 | 89 | 10 | 100 | 546 | 1000 | 764 |
-- – – – – – – – – -- – – – – – – – – -- – – – – – – – – -- -- -- - -
⬆ ⬆ ️
pivot high
low
low != high and pivot - low != k hence continue to next iteration
Third Iteration
k = 2
-- – – – – – – – – -- – – – – – – – – -- – – – – – – – – -- -- -- - --
| 1 | 2 | 3 | 67 | 89 | 10 | 100 | 300 | 546 | 1000 | 764 |
-- – – – – – – – – -- – – – – – – – – -- – – – – – – – – -- -- -- - --
⬆ ⬆ ⬆
low pivot high
low != high and pivot - low != k hence continue to next iteration
Fourth iteration (Solution)
k = 2
-- – – – – – – – – -- – – – – – – – – -- – – – – – – – – -- -- -- - --
| 1 | 2 | 3 | 10 | 89 | 67 | 100 | 300 | 546 | 1000 | 764 |
-- – – – – – – – – -- – – – – – – – – -- – – – – – – – – -- -- -- - --
⬆ ⬆ ⬆
low pivot high
pivot - low == k hence we found the solution. Return the element situated at index pivot
And now let's head back to implementation. The code is written in Swift which is the primary language I currently use.
func findKthSmallestElement(input: inout [Int], l: Int, r: Int, which: Int) -> Int {
if l == r {
return input[l]
}
if which >= 0 && which <= r - l {
let pivot = findPivot(input: &input, l: l, r: r)
if pivot - l == which {
return input[pivot]
} else {
if which < pivot - l {
return findKthSmallestElement(input: &input, l: l, r: pivot - 1, which: which)
} else {
return findKthSmallestElement(input: &input, l: pivot + 1, r: r, which: which - pivot - 1 + l)
}
}
}
return Int.max;
}
func findPivot(input: inout [Int], l: Int, r: Int) -> Int {
var min = l - 1
let pivot = input[r]
for i in l..<r {
if input[i] <= pivot {
min += 1
swapI(input: &input ,l: min, r: i)
}
}
min += 1
swapI(input: &input, l: min, r: r)
return min
}
To verify the output, let's write some unit tests,
func testFindSmallestKthElement() {
let array = [100, 300, 2, 3, 1000, 67, 89, 10, 764, 1, 546]
var input = array
XCTAssertEqual(findKthSmallestElement(input: &input, l: 0, r: input.count - 1, which: 0), 1)
input = array
XCTAssertEqual(findKthSmallestElement(input: &input, l: 0, r: input.count - 1, which: 3), 10)
input = array
XCTAssertEqual(findKthSmallestElement(input: &input, l: 0, r: input.count - 1, which: 5), 89)
input = array
XCTAssertEqual(findKthSmallestElement(input: &input, l: 0, r: input.count - 1, which: 7), 300)
input = array
// Request to find out of bounds kth minimum element
XCTAssertEqual(findKthSmallestElement(input: &input, l: 0, r: input.count - 1, which: 13), Int.max)
}
Time Complexity:
For best and average cases quick-select runs in linear time. Since algorithm keeps dividing array in half after every iteration, the size of input to inspect keep getting split in half. For array of size n, time complexity for average and best cases could be calculated with following formula,
n + (n/2) + (n/4) +..... 1 = O(2n) = O(n) using Infinite Geometric series formula.
For worst case scenarios, this algorithm works same way as quick sort algorithm leading to O(n^2) time complexity for the input of size n. This happens when the pivot element keep falling on extreme end and algorithm runs until it has only one element to process which is desired kth smallest element.
You can easily simulate this scenario by using sorted (in increasing order) array as an input asking algorithm to find 0th smallest element. The element is situated at the very first position, but ironically it still takes O(n^2) time to find the element.
Finding kth largest element
So this is fine, but how do we find kth largest element in the input array. It's exactly same except for a one line change in func findPivot function. The condition where we check if if input[i] <= pivot only needs to be changed to if input[i] >= pivot.
This works because earlier we were selectively sorting input array in increasing order so that we can grab kth smallest element from left. Now we still want to search from left, but instead array will be reverse sorted. So the place where we used to get kth smallest element, we will now be getting kth largest element.
Reference: GeeksForGeeks.com