Swiftならメモ化も最高にスッキリ書けます

Swiftでメモ化をスッキリと書く方法を紹介したいと思います。

この記事では、フィボナッチ数を例に説明します。
フィボナッチ数についてはWikipediaなど参照してください。
## メモ化を使わない例

その前にメモ化を使わないパターンです。

```swift:PureFibonacciFunc.swift
// n番目のフィボナッチ数を返す: 0, 1, 1, 2, 3, 5, 8, 13, ...
func fibonacci(_ n: Int) -> Int {
    return n < 2 ? n : fibonacci(n - 1) + fibonacci(n - 2)
}
```

素直な書き方ですが、上記の関数には_大きな欠点_があります。それは実行速度です。
たとえば手元のMacbook pro(13-inch, Mid 2014)で試してみたところ、`fibonacci(50)`を計算するのに__45秒__かかりました。

なぜこんなにも時間がかかるかというと、_無駄な計算_をしているからです。
少し詳しく説明します。

#### 実行速度が遅い理由

例えば`fibonacci(5)`を計算したとしましょう。すると、以下のようなコールグラフになります。![Fibs.png](https://qiita-image-store.s3.amazonaws.com/0/130238/bc8362e1-7261-2300-21b0-64bb63cc64be.png)

まず <font color="Orange">fibonacci(5)</font> を計算するのに <font color="#005B00">fibonacci(4)</font> と <font color="#019EB4">fibonacci(3)</font>  を計算することになります。<font color="#005B00">fibonacci(4)</font> を計算するために <font color="#019EB4">fibonacci(3)</font>, <font color="#0272C6">fibonacci(2)</font> が必要です。
これだけでも <font color="#019EB4">fibonacci(3)</font> を2回も計算しています。
さらに、コールグラフ全体を見ると <font color="#0272C6">fibonacci(2)</font> は3回も計算していることがわかります。

5番目のフィボナッチ数を計算するのにこれだけ無駄な計算が起きています。いわんや`fibonacci(50)`になると大きな非常に多くの無駄な計算をしていることは想像に難くないです。実行速度が落ちている理由はこれです。

上記の非効率を解決するために重要になるのがメモ化です。

## メモ化とは

上記のアルゴリズムで問題だったのは_同じ計算を繰り返していたから_でした。なので、一度計算したものはキャッシュしておいて、2回目以降はキャッシュした値を使うことにします。これが_メモ化_です。
概念的には下記のコールグラフのようになります。
![Fibs Memoize.png](https://qiita-image-store.s3.amazonaws.com/0/130238/0ccf3e6a-4c28-2c73-e2d6-e2eb8c177152.png)

最初に<font color="#005B00">fibonacci(4)</font> と <font color="#019EB4">fibonacci(3)</font> を計算するのは同じですが、__その結果をキャッシュ__しておき、2回目以降は計算を省くことで計算速度を速めます。コールグラフでは計算が省かれたものは薄く表示しています。

説明は以上にして、メモ化を実装していきましょう。

## メモ化の実装(素朴な方法)

```swift:MemoizedFibonacciFunc.swift
var fibonacciMemo = [Int: Int]() // この連想配列に計算結果をキャッシュしておく

// n番目のフィボナッチ数を返す: 0, 1, 1, 2, 3, 5, 8, 13, ...
func fibonacci(_ n: Int) -> Int {
    // まずはキャッシュがあるか確認
    if let result = fibonacciMemo[n] {
        return result // キャッシュしていた結果を返す
    }
    let result = n < 2 ? n : fibonacci(n - 1) + fibonacci(n - 2) // キャッシュがないので計算する
    fibonacciMemo[n] = result // 結果をキャッシュする
    return result // キャッシュして、return
}

fibonacci(50) // 12586269025
```

メモ化を使ったのが上記のコードです。
手元のMacでの実行時間は__0.01秒でした__!!劇的に速くなっています。

ですが、このコード、すごく_読みにくいです_。
実際の計算を行なっている部分(つまり一番重要な部分)は
`let result = n < 2 ? n : fibonacci(n - 1) + fibonacci(n - 2)`
ですが、それを見つけるのに時間がかかってしまいます。
しかも関数の__外__に連想配列を定義しておかなければなりません。
もう少し良い方法はないでしょうか。

## メモ化の実装(スッキリ書ける方法)
ここからがキモです。
下記の様にわかりやすくて読みやすい関数を目指します。

```swift:MemoizedFibonacciFunc2.swift
let fibonacci = memoize { fibonacci, n in
    n < 2 ? n : fibonacci(n - 1) + fibonacci(n - 2)
}
```

上記では`memoize`という関数をつかっています。これはSwiftの標準ライブラリの関数_ではありません。_
この`memoize`は自前の関数であり、これこそが、わかりやすくて読みやすいメモ化を実現してくれます。

### memorize part 1

関数`memoize`の_第一弾_です(第二弾まであります)。

```swift:memoize
/// memoize part 1
func memoize<T: Hashable, U>(body: @escaping(T) -> U) -> (T) -> U {
    var memo = [T: U]()
    return {
        if let q = memo[x] {
            return q
        }
        let r = body(x)
        memo[x] = r
        return r
    }
} 
```

少し複雑ですね。どのような関数になっているか解説します。

#### 型
ジェネリックな関数になっています。`T`は`Hashable`に準拠していなければなりません。`U`は任意の型でOKです。
`T`が`Hashable`でないといけない理由は、連想配列のキーに使うためです。
#### 引数
`memoize`の引数は`body`です。型は`(T) -> U`です。つまり、Tを引数に取りUを返すクロージャ式です。
#### 戻り値
戻り値も、引数とおなじ`(T) -> U`です。
#### 関数の中身
さて、関数の中身をみていきましょう。
最初に宣言されているのは`memo`という連想配列です。この配列に計算結果をキャッシュします。
次にクロージャ式を返しています。クロージャの内部はシンプルで、`memo`に値がキャッシュされていればそれを使い、キャッシュがなければ計算してキャッシュに保存した後、計算結果をreturnしています。
ポイントは__クロージャ式が`memo`をキャプチャしている__ところです。_キャプチャのおかげで、関数の外に連想配列を定義する必要がなくなりました。_
#### 使い方

`memoize`の使い方をみてみます。
まずは素数判定関数で試してみます。

```swift:MemoizeUsage1
func isPrime(_ n: Int) -> Bool {
    for i in 1..<n {
        if n % i == 0 {
            return false
        }
    }
    return true
}

let isPrimeMemoize = memoize { n in
    isPrime(n)
}

isPrimeMemoize(1187) // true
isPrimeMemoize(1187) // true(メモ化した結果を返してくれる)
```

OKです！素数か否かを判定する`isPrime(_:)`の結果をメモ化してくれました。
(素数判定アルゴリズムが非常に非効率なものですがこの記事の主題ではないのでツッコミはナシでお願いします!)

フィボナッチ数でも試してみましょう。

```swift:MemoizeUsageError
// error: variable used within its own initial value
let fibonacci = memoize { n in
    n < 2 ? n : fibonacci(n - 1) + fibonacci(n - 2)
}
```

上記コードではエラーになってしまいます。
エラー文の通り、`fibonacci`の初期化のために`fibonacci`それ自体を使っているからです。こうしないといけないのは`fibonacci`が再帰的な関数だからです。
そこで、一旦nilで初期化しておきましょう。

```swift:MemoizeUsage2
// OK:
var fibonacci: ((Int) -> Int)! // 一旦、nilで初期化しておく
fibonacci = memoize { n in
    n < 2 ? n : fibonacci(n - 1) + fibonacci(n - 2)
}

fibonacci(6) // 8
fibonacci(8) // 21
```

OKです。動きました。
しかし、まだ改善の余地があります。わざわざnilで初期化しないといけないし、なにより`fibonacci`がmutableになってしまっています。

なんとか改善の方法はないでしょうか。

### memoize part 2

memoizeの第二弾です。これが完成形です。

```swift:MemoizeAlt
func memoize<T: Hashable, U>(body: @escaping((T) -> U, T) -> U) -> (T) -> U {
    var memo = Dictionary<T, U>()
    var result: ((T) -> U)! // nilでの初期化はmemoizeの内部で行う
    result = { x in
        if let q = memo[x] {
            return q
        }
        // `body`の第一引数に`result`を渡す
        let r = body(result, x)
        memo[x] = r
        return r
    }
    return result // 非オプショナル型として返す
}
```

第一弾との違いは、`body`の型です。クロージャ式なのは同じですが、クロージャ式の引数の型が`T`から`(T) -> U, T`になっています。つまり新たな引数`(T) -> U`が加えられています。
もう一つの違いは、クロージャ式の初期化は`memoize`の内部で行っていることです。つまり、第一弾でネックになっていたことを`memoize`の内部に押し込んでいます。

使い方は以下です。

```swift:MemoizeAltUsage
let fibonacci = memoize { fibonacci, n in
    // body内部の`fibonacci`は第一引数の`fibonacci`
    n < 2 ? n : fibonacci(n - 1) + fibonacci(n - 2)
}

fibonacci(6) // 8 🎉
fibonacci(8) // 21 🎉
```

新しく`body`に追加された引数を再帰に使っています。このおかげで`fibonacci`の初期化のために`fibonacci`それ自体を使うことを回避できます。ただし、実際は、`let`で宣言している`fibonacci`と`body`の引数の`fibonacci`は同じ関数を参照しています。
関数が参照型であることを上手く使っています。

## その他の計算例

### 階乗

メモ化の応用例としてよく使われる階乗の計算の利用方法を試してみます。

```swift:MemoizeFactorial
let factorial = memoize { factorial, n in
    n < 2 ? n : n * factorial(n - 1)
}

factorial(5) // 120
factorial(7) // 5040
```

### 素数判定

`isPrime(_:)`は再帰的な構造ではないので第二引数のみ使います。

```swift:MemoizePrime
let isPrimeMemoized = memoize { _, n in
    return isPrime(n)
}

isPrimeMemoized(10) // false
```

## 参考にした資料

2014年のWWDCのAdvanced Swiftというセッションを参考にしました。
https://developer.apple.com/videos/play/wwdc2014/404/