Quickselect algorithm in Swift

Question

I recently answered a question on Stack Overflow about finding the k-largest element in an array, and present my implementation of the Quickselect algorithm in Swift for review.

This is essentially the iterative version described in Wikipedia: Quickselect, only that the paritioning does not move the pivot element to the front of the second partition. That saves some swap operations but requires an additional check when updating the lower bound.

The language is Swift 3, compiled with Xcode 8.1.

extension Array where Element: Comparable {

    public func kSmallest(_ k: Int) -> Element {
        precondition(1 <= k && k <= count, "k must be in the range 1...count")

        var a = self // A mutable copy.
        var low = startIndex
        var high = endIndex

        while high - low > 1 {

            // Choose random pivot element:
            let pivotElement = a[low + Int(arc4random_uniform(UInt32(high - low)))]

            // Partition elements such that:
            //   a[i] <  pivotElement    for low <= i < pivotIndex,
            //   a[i] >= pivotElement    for pivotIndex <= i < high.
            var pivotIndex = low
            while a[pivotIndex] < pivotElement {
                pivotIndex += 1
            }
            for i in pivotIndex+1 ..< high {
                if a[i] < pivotElement {
                    swap(&a[pivotIndex], &a[i])
                    pivotIndex += 1
                }
            }

            if k <= pivotIndex {
                // k-smallest element is in the first partition:
                high = pivotIndex
            } else if k == pivotIndex + 1 {
                // Pivot element is the k-smallest:
                return pivotElement
            } else {
                // k-smallest element is in the second partition
                low = pivotIndex
                if a[low] == pivotElement {
                    low += 1
                }
            }
        }

        // Only single candidate left:
        return a[low]
    }

    public func kLargest(_ k: Int) -> Element {
        return kSmallest(count + 1 - k)
    }
}

Examples:

let a = [ 9, 7, 6, 3, 4, 2, 5, 1, 8 ]
for k in 1 ... a.count {
    print(a.kSmallest(k))
}
// 1 2 3 4 5 6 7 8 9

let b = [ "b", "a", "c", "h" ]
print(b.kLargest(2))
// "c"

Feedback on all aspects of the code is welcome, such as (but not limited to):

• Possible improvements (speed, clarity, swiftyness, ...).
• Naming. In particular: what would be a better name for kSmallest(_ k: Int) in consideration of the Swift API Design Guidelines?
• The check if a[low] == pivotElement looks artificial, but without that an infinite loop can occur, e.g. for an array with all elements equal. Is there a better solution?

---

Remark: To make this code compile with Swift 4 (or later), just replace the line

swap(&a[pivotIndex], &a[i])

with

a.swapAt(pivotIndex, i)

Answer 1

Let's start by updating this line to Swift 4.2 :

let pivotElement = a[Int.random(in: low..<high)]

In my tests, Int.random(in:) is way faster than Int(arc4random_uniform), and the comparisons won't take that into consideration.

---

Efficiency

1st improvement

There is a small improvement to this algorithm, but it still gives a performance gain by rearranging the conditions from the most probable to the least, in order to take advantage of shortcut execution:

if k <= pivotIndex {
    // k-smallest element is in the first partition:
    high = pivotIndex
} else if k > pivotIndex + 1 {
    // k-smallest element is in the second partition
    low = pivotIndex
    if a[low] == pivotElement {
        low += 1
    }
} else {
    // Pivot element is the k-smallest:
    return pivotElement
}

The first two conditions are equiprobable. The least probable case is left for last.

Benchmarks

The benchmarking code is the following:

let a = Array(1...10_000).shuffled()

var timings: [Double] = []
for k in 1 ... a.count {
    let start = mach_absolute_time()
    _ = a.kSmallest(k)
    let end = mach_absolute_time()
    timings.append(Double(end - start)/Double(1e3))
}
let average: Double = timings.reduce(0, +) / Double(timings.count)
print(average, "us")

var timings2: [Double] = []
for k in 1 ... a.count {
    let start = mach_absolute_time()
    _ = a.kSmallest2(k)
    let end = mach_absolute_time()
    timings2.append(Double(end - start)/Double(1e3))
}
let average2: Double = timings2.reduce(0, +) / Double(timings2.count)
print(average2, "us")

It prints the average time for looking up one kth smallest element.

kSmallest is the original, kSmallest2 is the new one. They both operate on the same array a to ensure fairness.

kSmallest2 is up to 7μs faster per lookup. The fluctuation is due to the randomness of the arrangement of the elements of the array. Which translates into up to ~70ms execution time gain for a 10.000-element array:

kSmallest    1.215636265 s (total time)
kSmallest2   1.138085315 s (total time)

In the worst case, in my tests, kSmallest2 may rarely be 2μs slower per lookup, and it is to be blamed on the randomness of choosing a pivot. Comparisons should probabilistically favor the second version.

---

2nd improvement

The following improvement concerns arrays with duplicates, and avoids unnecessary loops:

while a[low] == pivotElement, k - low > 1 {
    low += 1
}

Instead of hopping by one index alone:

if a[low] == pivotElement {
    low += 1
}

Benchmarks

The following code was used:

//As suggested by Tim Vermeulen
let a = (0..<100).flatMap { Array(repeating: $0, count: Int.random(in: 10..<30)) }
    .shuffled()

var timings1: [Double] = []
for k in 1 ... a.count {
    let start = mach_absolute_time()
    _ = a.kSmallest(k)
    let end = mach_absolute_time()
    timings1.append(Double(end - start)/Double(1e6))
}
let average1: Double = timings1.reduce(0, +) / Double(timings1.count)
print("kSmallest", average1, "ms")

var timings2: [Double] = []
for k in 1 ... a.count {
    let start = mach_absolute_time()
    _ = a.kSmallest2(k)
    let end = mach_absolute_time()
    timings2.append(Double(end - start)/Double(1e6))
}
let average2: Double = timings2.reduce(0, +) / Double(timings2.count)
print("kSmallest2", average2, "ms")


var timings3: [Double] = []
for k in 1 ... a.count {
    let start = mach_absolute_time()
    _ = a.kSmallest3(k)
    let end = mach_absolute_time()
    timings3.append(Double(end - start)/Double(1e6))
}
let average3: Double = timings3.reduce(0, +) / Double(timings3.count)
print("kSmallest3", average3, "ms")

kSmallest3 has both the 1st and 2nd improvements.

Here are the results:

kSmallest  0.0272 ms
kSmallest2 0.0267 ms
kSmallest3 0.0236 ms

In an array with a high number of duplicates, the original code is now ~13% faster by implementing both improvements. That percentage will grow with the richness in duplicates, and a higher array count. If the array has unique elements, kSmallest2 is naturally the fastest since it'll be avoiding unnecessary checks.

---

3rd improvement (a 🐞 fix?)

There are unnecessary loops where the the random index is that of a pivot element which is already in its rightful place/order. These elements aren't swapped by the code, since they are already well-placed. These elements are the ones that fall in the case of k > pivotIndex + 1 and the low index is equal to pivotIndex. An endless loop may occur if Int.random(in: low..<high) always returns low + 1.

The following code prevents such an (admittedly unlikely) endless loop:

var orderedLows: Set<Int> = []  //This will contain the indexes of elements that are already well ordered
while high - low > 1 {
    // Choose random pivot element:
    var randomIndex: Int
    repeat {
        randomIndex = Int.random(in: low..<high)
    } while orderedLows.contains(randomIndex) &&
    !orderedLows.isSuperset(of: Array<Int>(low..<high))

    let pivotElement = a[randomIndex]

    ...

    } else if k > pivotIndex + 1 {
        // k-smallest element is in the second partition
        if low == pivotIndex
        {
            orderedLows.insert(randomIndex)
        }
        low = pivotIndex
        while a[low] == pivotElement, k - low > 1 {
            low += 1
        }
    }

Benchmarks

The following array was operated on with the same benchmarking code as in the second improvment:

let a = (0..<100).flatMap { Array(repeating: $0, count: Int.random(in: 10..<30)) }
    .shuffled()

And here are the results:

kSmallest  0.0662 ms
kSmallest2 0.0639 ms
kSmallest4 0.0575 ms

kSmallest4 being the version with all three improvements, and it is faster than all previous versions. The a array was purposefully chosen to be rich in duplicates to heighten the possibilty of elements that are already in the correct order. If not, kSmallest4 doesn't show any flagrant improvement.

---

Naming

1) pivotIndex is confusing, one would expect it t be the index of pivotElement, but it's not.

2) a isn't very descriptive. Its name neither conveys its mutability nor that it is a copy of the initial array.

3) Such an algorithm, née FIND, is commonly known as quickSelect (more names could be found here). Personally, I would prefer nthSmallest(_ n: Int), for the following reasons:

• n instead of k since the latter is usually used in constant naming
• nth instead of n denotes ordinality
• Since the comparison predicate isn't provided, nthSmallest(_ n: Int) would be preferable to nthElement(_ n: Int) since it means that we'll be comparing elements in an ascending order.

---

Readability/Alternative approach

If readability and conciseness are paramount, here is an alternative that uses the partition(by:) method applied to the array slice mutableArrayCopy[low..<high] :

public func nthSmallest(_ n: Int) -> Element {
    precondition(count > 0, "No elements to choose from")
    precondition(1 <= n && n <= count, "n must be in the range 1...count")

    var low = startIndex
    var high = endIndex
    var mutableArrayCopy = self

    while high - low > 1 {
        let randomIndex = Int.random(in: low..<high)
        let randomElement = mutableArrayCopy[randomIndex]
        //pivot will be the index returned by partition
        let pivot = mutableArrayCopy[low..<high].partition { $0 >= randomElement }
        if n < pivot + 1 {
            high = pivot
        } else if n > pivot + 1 {
            low = pivot
            //Avoids infinite loops when an array has duplicates
            while mutableArrayCopy[low] == randomElement, n - low > 1 {
                low += 1
            }
        } else {
            return randomElement
        }
    }
    // Only single candidate left:
    return mutableArrayCopy[low]
}

In my tests, the more duplicates in the array, the more this version gets on a par (if not better) with the original code, but a bit slower when the array has unique elements.

Here are some tests:

[1, 3, 2, 4, 7, 8, 5, 6, 9, 10].nthSmallest(6)          //6

[10, 20, 30, 40, 50, 60, 20, 30, 20, 10].nthSmallest(4) //20

["a", "a", "a", "a", "a"].nthSmallest(3)                //"a"

A recursive rewrite of the original code seems more readable. But at the cost of execution time, since each time the function is called, a mutable copy of the array will be created.

Answer 2

There are a lot of loops in here. I would probably turn at least one of them into a recursive function; actually, refactoring high low and pivotIndex into let constants with recursion should be faster due to optimizations. (I'm not that good of an FP programmer to mock something up for you quickly).

Otherwise, one suggestion I'd make (just my preference), you have two major loops smooshed next to each other, perhaps another empty line could help readability:

            // Partition elements such that:
            //   a[i] <  pivotElement    for low <= i < pivotIndex,
            //   a[i] >= pivotElement    for pivotIndex <= i < high.
            var pivotIndex = low

            while a[pivotIndex] < pivotElement {
                pivotIndex += 1
            }

            for i in pivotIndex+1 ..< high {
                if a[i] < pivotElement {
                    swap(&a[pivotIndex], &a[i])
                    pivotIndex += 1
                }
            }

I would definitely try refactoring some of this to FP, then running it through some measureBlock()s in Xcode to see if it's faster.

Hope this helps!