partition.swift

extension MutableCollection {
    /// Reorders the elements of the collection such that all the elements 
    /// that match the given predicate are after all the elements that do 
    /// not match the predicate.
    ///
    /// After partitioning a collection, there is a pivot index `p` where 
    /// no element before `p` satisfies the `belongsInSecondPartition` 
    /// predicate and every element at or after `p` satisfies 
    /// `belongsInSecondPartition`.
    /// 
    /// In the following example, an array of numbers is partitioned by a
    /// predicate that matches elements greater than 30.
    ///
    ///     var numbers = [30, 40, 20, 30, 30, 60, 10]
    ///     let p = numbers.partition(by: { $0 > 30 })
    ///     // p == 5
    ///     // numbers == [30, 10, 20, 30, 30, 60, 40]
    ///
    /// The `numbers` array is now arranged in two partitions. The first 
    /// partition, `numbers.prefix(upTo: p)`, is made up of the elements that 
    /// are not greater than 30. The second partition, `numbers.suffix(from: p)`, 
    /// is made up of the elements that *are* greater than 30.
    ///
    ///     let first = numbers.prefix(upTo: p)
    ///     // first == [30, 10, 20, 30, 30]
    ///     let second = numbers.suffix(from: p)
    ///     // second == [60, 40]
    ///
    /// - Parameter belongsInSecondPartition: A predicate used to partition
    ///   the collection. All elements satisfying this predicate are ordered 
    ///   after all elements not satisfying it.
    /// - Returns: The index of the first element in the reordered collection
    ///   that matches `belongsInSecondPartition`. If no elements in the
    ///   collection match `belongsInSecondPartition`, the returned index is
    ///   equal to the collection's `endIndex`.
    ///
    /// - Complexity: O(n)
    public mutating func partition(
        by belongsInSecondPartition: @noescape (Iterator.Element) throws -> Bool
    ) rethrows -> Index {
        var pivot = startIndex
        while true {
            if pivot == endIndex {
                return pivot
            }
            if try belongsInSecondPartition(self[pivot]) {
                break
            }
            formIndex(after: &pivot)
        }
        
        var i = index(after: pivot)
        while i < endIndex {
            if try !belongsInSecondPartition(self[i]) {
                swap(&self[i], &self[pivot])
                formIndex(after: &pivot)
            }
            formIndex(after: &i)
        }
        return pivot
    }
}

extension MutableCollection where Self: BidirectionalCollection {
    /// Reorders the elements of the collection such that all the elements 
    /// that match the given predicate are after all the elements that do 
    /// not match the predicate.
    ///
    /// After partitioning a collection, there is a pivot index `p` where 
    /// no element before `p` satisfies the `belongsInSecondPartition` 
    /// predicate and every element at or after `p` satisfies 
    /// `belongsInSecondPartition`.
    /// 
    /// In the following example, an array of numbers is partitioned by a
    /// predicate that matches elements greater than 30.
    ///
    ///     var numbers = [30, 40, 20, 30, 30, 60, 10]
    ///     let p = numbers.partition(by: { $0 > 30 })
    ///     // p == 5
    ///     // numbers == [30, 10, 20, 30, 30, 60, 40]
    ///
    /// The `numbers` array is now arranged in two partitions. The first 
    /// partition, `numbers.prefix(upTo: p)`, is made up of the elements that 
    /// are not greater than 30. The second partition, `numbers.suffix(from: p)`, 
    /// is made up of the elements that *are* greater than 30.
    ///
    ///     let first = numbers.prefix(upTo: p)
    ///     // first == [30, 10, 20, 30, 30]
    ///     let second = numbers.suffix(from: p)
    ///     // second == [60, 40]
    ///
    /// - Parameter belongsInSecondPartition: A predicate used to partition
    ///   the collection. All elements satisfying this predicate are ordered 
    ///   after all elements not satisfying it.
    /// - Returns: The index of the first element in the reordered collection
    ///   that matches `belongsInSecondPartition`. If no elements in the
    ///   collection match `belongsInSecondPartition`, the returned index is
    ///   equal to the collection's `endIndex`.
    ///
    /// - Complexity: O(n)
    public mutating func partition(
        by belongsInSecondPartition: @noescape (Iterator.Element) throws -> Bool
    ) rethrows -> Index {
        var lo = startIndex
        var hi = endIndex
        
        // 'Loop' invariants (at start of Loop, all are true):
        // * lo < hi
        // * predicate(self[i]) == false, for i in startIndex ..< lo
        // * predicate(self[i]) == true, for i in hi ..< endIndex
        
        Loop: while true {
            FindLo: repeat {
                while lo < hi {
                    if try belongsInSecondPartition(self[lo]) { break FindLo }
                    formIndex(after: &lo)
                }
                break Loop
            } while false
            
            FindHi: repeat {
                formIndex(before: &hi)
                while lo < hi {
                    if try !belongsInSecondPartition(self[hi]) { break FindHi }
                    formIndex(before: &hi)
                }
                break Loop
            } while false
            
            swap(&self[lo], &self[hi])
            formIndex(after: &lo)
        }
        
        return lo
    }
}

var n = [30, 40, 20, 30, 30, 60, 10]
let p = n.partition(by: { $0 > 30 })
// n == [30, 10, 20, 30, 30, 60, 40]
// p == 5