[Swift] 条件を指定して配列を分割できる`partition(by:)`の使い方

partition(by:)とは

与えた条件に一致するすべての要素が、一致しないすべての要素の後に来るように、コレクションの要素を並べ替えます。

使い方

var numbers = Array(0...10)
// [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

let p = numbers.partition(by: { $0 % 2 == 0 })
// 5

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

let odds = numbers[..<p] // [9, 1, 7, 3, 5]
let evens = numbers[p...] // [4, 6, 2, 8, 0, 10]

注意点

• partition(by:)は条件に一致する要素が始まるindexを返すのであって、新しい配列を返すわけではない。
  • つまりpartition(by:)は自身を変化させるmutating func (元の配列を宣言する時はletではなくvar)
• 要素を並び替えた後、条件を満たさない前半(or満たす後半)の要素の連続の中で元の順番は保証されていない

問題点 & 対応策

前述したように、partition(by:)は条件によって配列を２つに分割することができますが、元の配列の順番が重要で、分割後もそれを維持したいケースにこのメソッドは向いていないです(そして多分そのケースの方が多い🤫)

対応策として

Arrayを拡張する

extension Array {
    /// 条件を満たさない要素の配列と、満たす配列をtupleにして返す
    func partitioned(predicate: (Element) -> Bool) -> (unsatisfied: [Element], satisfied: [Element]) {
        return reduce(into: ([Element](), [Element]())) {
            if predicate($1) {
                $0.1.append($1)
            } else {
                $0.0.append($1)
            }
        }
    }
}

使い方

let partitioned = numbers.partitioned(predicate: { $0 % 2 == 0})

print(partitioned.unsatisfied) // [1, 3, 5, 7, 9]
print(partitioned.satisfied) // [0, 2, 4, 6, 8, 10]

// 順番はそのまま！

ちなみに、今回extensionに利用したreduce(into:_:)については、こちらに素晴らしい記事があります。

また、パフォーマンスについてはreduce(into:)ではなく素直にfor inを使って回す方法もありますが、リリースビルドにおいては問題無いようです。
(あっても自分はreduceで書きたいが🤐)

最後に

当然ですが、分割した配列内の要素の順番を気にしないのであれば、パフォーマンスの観点からSwift標準のpartition(by:)を使うべきですね😌