SwiftでArrayの要素を振り分けたい

Advent Calendar 書くの遅れてゴメンなさい(´・ω・`)

SwiftでRubyのpartitionメソッドみたいな、
条件を満たす配列と満たさない配列に分割する関数が欲しかったので実装しました。
partition/Module: Enumerable (Ruby 2.2.3)

ankurp/Dollar.swiftのpartitionは分割数やステップ数を細く決められるが、
今回シンプルに条件で二分割したいだけ。

implement

extensionでArrayに直接追加します。

Rubyのpartitionは[[真の要素...], [偽の要素...]]を二重配列で返していますが、
今回はScalaのpartitionに習ってtupleで返却するようにします。
Scalaコレクションメソッドメモ(Hishidama's Scala collection method Memo)

また、Swiftのtupleは「要素に名前を付けられる + パターンマッチで取り扱いやすい¹」
ので配列ではなくtupleを採用しました。

以下、実装です。

extension Array {
    public func partition(@noescape predicate: Element -> Bool) -> (include: [Element], exclude: [Element]) {
        return reduce(([], [])) { (acc, x) in
            if predicate(x) {
                return (acc.0 + [x], acc.1)
            } else {
                return (acc.0, acc.1 + [x])
            }
        }
    }
}

クロージャを受け取って即時評価するのが決まっているので、
引数で受け取る条件関数predicateにnoescapeアノテーションを指定しています。
Swiftの @noescape をもっと使おう - Qiita

すぐに評価することが決まっている関数の場合はどうでしょう。そういうときはキャプチャする必要が無いし、self.も不要です。キャプチャしない方がパフォーマンス上でも有利でしょうし、最適化しやすいかもしれません。このようにすぐに評価され> ることがわかっているクロージャのために追加された新しいアノテーションがnoescapeです。
Swift 1.2 - cockscomblog?

sample

let isEven: Int -> Bool = { $0 % 2 == 0 }

Array(1...5).partition(isEven) // ([2, 4], [1, 3, 5])

戻り値のtupleの要素に名前を定義しているので、名前指定で直感的に取り出せます。
またtupleのパターンマッチングで分割代入も行えるので便利!!

let result = Array(1...5).partition(isEven)
result.includes // [2, 4]
result.excludes // [1, 3, 5]

// パターンマッチで簡単に取り出せる
let (trueAry, falseAry) = Array(1...5).partition(isEven)
trueAry // [2, 4]
falseAry // [1, 3, 5]

おまけ

partition関数を使ってクイックソートを書いてみます。
++ 配列のhead/tail分割用のcomputed-properties decompose²をArrayに追加しています。

extension Array {
    // 先頭とそれ以外に分割したtuple?を返す
    public var decompose: (head: Element, tail: [Element])? {
        return (count > 0) ? (self[0], Array(self[1..<count])) : nil
    }
}

func quickSort(ns: [Int]) -> [Int] {
    if let (pivot, rest) = ns.decompose {
        let (l, g) = rest.partition { $0 < pivot }
        return quickSort(l) + [pivot] + quickSort(g)
    } else {
        return []
    }
}

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

すっきりと書けた。
tupleのパターンマッチ便利ですね!

1. Swiftでhead、tailにパターンマッチできる遅延リストの、List型のhead/tailパターンマッチも考えましたが、今回の用途ではtupleが適切かと。 ↩

2. Functional Snippet #1: Decomposing Arrays · objc.io ↩