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?