Quickselect algorithm in Swift 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) -&gt; Element {
        precondition(1 &lt;= k &amp;&amp; k &lt;= count, "k must be in the range 1...count")

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

        while high - low &gt; 1 {

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

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

            if k &lt;= 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) -&gt; 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(&amp;a[pivotIndex], &amp;a[i])


with

a.swapAt(pivotIndex, i)
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] &lt;  pivotElement    for low &lt;= i &lt; pivotIndex,
            //   a[i] &gt;= pivotElement    for pivotIndex &lt;= i &lt; high.
            var pivotIndex = low

            while a[pivotIndex] &lt; pivotElement {
                pivotIndex += 1
            }

            for i in pivotIndex+1 ..&lt; high {
                if a[i] &lt; pivotElement {
                    swap(&amp;a[pivotIndex], &amp;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!