Swiftで要素の順番を保持しつつ2つの配列の差分を取る

久々の更新です。

~~会社のTech Blogを書くのも溜まってるから頑張らなきゃ…~~

今日は、2つの配列の差分を取得するextensionと、取得するときに要素の並びを保持するにはどうしたらいいかについて触れてみます。

配列の差分

今回想定する差分というのは、具体的には xor の処理になります。

つまり、 2つの配列を比べて、重複している要素を除いたもの を取得する形になります。

例

let A = [1,2,3,4]
let B = [2,3,5]
let result = A.diff(B)
print(result) // [1,4,5]

この例では、A,B両方に含まれている 2, 3 は除外され、 1, 4, 5 が取得できます。

実装してみる

Swiftには Set が用意されていて、2つのSetのxorを取る関数 exclusiveOr も用意されているので、それを使ってみます。
ただし、要素が Hashable である必要があります。

extension CollectionType where Generator.Element: Hashable {
    typealias E = Generator.Element
    func diff(other: [E]) -> [E] {
        return Array((Set(self).exclusiveOr(Set(other))))
    }
}

こんな形でCollectionTypeのextensionとして定義してみます。
一度配列を Set に包んで、xor を取って再度配列に戻します。
実行してみます。

let a = [1, 2, 3, 4]
let b = [8, 2, 3, 9]
_ = a.diff(b) // => [9, 4, 1, 8]

確かに、共通で存在する 2, 3 は削除されました。が、あれ…？

ここで問題点が…

よく見てみると、結果で得られる配列の要素の順番が ぐちゃぐちゃ になっています。
できれば、 a.diff(b) としたので、 a+b と結合したときの順序と同じにしたいですね。
つまり、今回のケースでは、 [1, 4, 8, 9] となって欲しいわけです。

順序を保ちつつ差分を取ってみる

順序を保つには、Setを使わずに、配列のまま扱う必要があります。
実装する方法は色々あると思いますが、今回は以下のアプローチで挑もうと思います。

前提として、要素が Hashable に適合しているものとする(一意性がある)
まず、配列 A と B を足した配列 all を作る
各要素の出現数を、要素を key 出現数を value として格納する連想配列を用意する
配列 all の要素を見ていき、出現数を記録していく

ということで、実装してみます。

extension CollectionType where Generator.Element: Hashable {
    typealias E = Generator.Element
    func diff(other: [E]) -> [E] {
        let all = self + other
        var counter: [E: Int] = [:]
        all.forEach { counter[$0] = (counter[$0] ?? 0) + 1 }
        return all.filter { (counter[$0] ?? 0) == 1 }
    }
}

これで、しっかりと順序を保って xor を取ることができます。

let a = [1, 2, 3, 4]
let b = [8, 2, 3, 9]
_ = a.diff(b) // => [1, 4, 8, 9]
_ = b.diff(a) // => [8, 9, 1, 4]

番外: filterでなくて 2つの`flatMap`を使ってみる

先ほどのコードで、

        return all.filter { (counter[$0] ?? 0) == 1 }

としている箇所があったのですが、ここをflatMapを使ってみると、こんな感じになります。

extension CollectionType where Generator.Element: Hashable {
    typealias E = Generator.Element
    func diff(other: [E]) -> [E] {
        let all = self + other
        var counter: [E: Int] = [:]
        all.forEach { counter[$0] = (counter[$0] ?? 0) + 1 }
        return all.flatMap { e in counter[e].flatMap { $0 == 1 ? e : nil } }
    }
}

2つ、flatMapがでてくるのですが、前者はArrayに対してのflatMapで、後者は Optional に対してのflatMapになっています。
前者の方は、要素がnilだったらその要素を取り除くのに用います。
後者の方は、値が None だったら何もせず nil を返し、 Some だったら、 closure を適応して、適応した値を Optional で返します。
今回のケースだと、filterで書いたほうが綺麗に書けますが、flatMapで、 ?? を使わずに書いてみるのもいいのかなって思いました。

余談

来週あたりに、会社のTech Blog書きます。書く書く詐欺にならないように頑張ります。まる。