Cfrtjryk

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)

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

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.

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!