Swift で配列の直積（Cartesian product）を作る

はじめに

Swift で２つの配列から直積を作ることができれば、ループの入れ子をなくすことができます。

before

[1,2,3].forEach { i in
    ["a", "b", "c"].forEach { s in
        print(i, s)
    }
}

after

[1,2,3].product(["a", "b", "c"]).forEach { (i, s) in
    print(i, s)
}

こんな感じでコードを書きたかったので、Swift での実装方法を調べてみました。

環境

• Apple Swift version 5.3 (swiftlang-1200.0.29.2 clang-1200.0.30.1)

直積（Cartesian product）を作る

Array の直積

２つの配列からそれぞれの要素を取り出して、すべての組み合わせを要素に持つ配列を作ります。

２つの配列

let array1 = [1, 2, 3]
let array2 = ["foo", "bar", "baz"]

変換後

[(1, "foo"), (1, "bar"), (1, "baz"), (2, "foo"), (2, "bar"), (2, "baz"), (3, "foo"), (3, "bar"), (3, "baz")]

array1 の各要素ごとに array2 の各要素と結合したタプルのリストを作ります。
二次元配列を一次元配列に変換するために flatMap を使用します。

２つの配列から直積を作るコード

let product = array1.flatMap { x in
    array2.map { y in
        (x, y)
    }
}

LazySequence を使った高速化

上記のコードは product を生成するタイミングで array1 と array2 に対する二重ループ処理が実行されます。これは、配列の要素数が大きいほどオブジェクトの生成に時間がかかることを意味しています。

配列要素を参照するまでループ処理を遅延させることができれば、オブジェクトの生成を高速化することができます。

Swift では Array をはじめとする Sequence に準拠した型には lazy プロパティが提供されており、要素の参照を遅延させる事ができます。

遅延処理で直積を作るコード

let product = array1.lazy.flatMap { x in
    array2.lazy.map { y in
        (x, y)
    }
}

このように記述することで、実際に要素を参照するタイミングまで二重ループ処理を遅延させることができます。

forEach の呼び出しで二重ループ処理が実行される

product.forEach { item in
    print(item)
}

first の呼び出しで二重ループ処理は最初の１回で停止する

product.first

Array vs. LazySequence

とはいえ、何でも LazySequence で遅延実行すれば良いかというとそんなこともありません。

Array の特定の要素への参照が $O(1)$ なのに対して、LazySequence は $O(N)$ なので、実際のところ Array にするか LazySequence するかは一長一短です。

	オブジェクト生成時のコスト	特定の要素への参照時のコスト
Array	$O(N)$	$O(1)$
LazySequence	$O(1)$	$O(N)$

処理の共通化

特定の配列に対する直積の生成が理解できたところで、今度は任意の配列に対して直積を生成するべく処理の共通化を試みます。

Array

まずは Array の直積を作ってみます。

extension Sequence {
    func product<T: Sequence>(_ other: T) -> [(Element, T.Element)] {
        flatMap { x in
            other.map { y in
                (x, y)
            }
        }
    }
}

Sequence を拡張することで、配列をはじめとする Sequence に準拠した型すべてに対して、直積の Array を作ることができるようになります。

// [(1, true), (1, false), (2, true), (2, false), (3, true), (3, false)]
[1, 2, 3].product([true, false])

LazySequence

LazySequence を関数の戻り値にする場合は、どのような型を指定すれば良いでしょうか？

extension Sequence {
    func product<T: Sequence>(_ other: T) -> ??? { // ← 型は何？
        lazy.flatMap { x in
            other.lazy.map { y in
                (x, y)
            }
        }
    }
}

この場合の戻り値の型は次のようになります（とても長い！！）

LazySequence<FlattenSequence<LazyMapSequence<LazySequence<Self>.Elements, LazyMapSequence<LazySequence<T>.Elements, (Self.Element, LazySequence<T>.Element)>>>>

Swift では、lazy で遅延処理するために、flatMap や map などの演算処理を、FlattenSequence、LazyMapSequence のような型情報として保持する必要があるみたいで、演算が増えるほどに、型のネストが深くなります。

利用する側としては、内部で実行する演算処理について関心は無く、２つの組のタプルを要素にしたシーケンスであるということだけ知っていれば良いはずです。

このような場合は、AnySequence でラップすると分かりやすいです。

extension Sequence {
    func product<T: Sequence>(_ other: T) -> AnySequence<(Element, T.Element)> {
        AnySequence(lazy.flatMap { x in
            other.lazy.map { y in
                (x, y)
            }
        })
    }
}

こうすることで、要素が (Element, T.Element) で Sequence に準拠した型であることを明示することができます。メソッドの戻り値の型から遅延処理されることが隠蔽されてしまったことは難点ですが、実際には遅延処理されます。

AnySequence でラップすることで可読性はアップするものの、真の型である LazySequence が隠蔽されてしまう点には一長一短があります。

たとえば、Sequence やそれに準拠している AnySequence には first(where:) しかなく、LazySequence は first(where:) に加えて first プロパティもあります。

// first のかわりに first(where:) を利用する
[1,2,3].product([true, false]).first { _ in true }

ほとんどの場合で AnySequence でラップして返すで問題はなさそうですが、どちらが適切かはケースバイケースで考える必要がありそうです。