Swiftでメモ化をスッキリと書く方法を紹介したいと思います。
この記事では、フィボナッチ数を例に説明します。
フィボナッチ数についてはWikipediaなど参照してください。
メモ化を使わない例
その前にメモ化を使わないパターンです。
// 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)を計算したとしましょう。すると、以下のようなコールグラフになります。
まず fibonacci(5) を計算するのに fibonacci(4) と fibonacci(3) を計算することになります。fibonacci(4) を計算するために fibonacci(3), fibonacci(2) が必要です。
これだけでも fibonacci(3) を2回も計算しています。
さらに、コールグラフ全体を見ると fibonacci(2) は3回も計算していることがわかります。
5番目のフィボナッチ数を計算するのにこれだけ無駄な計算が起きています。いわんやfibonacci(50)になると大きな非常に多くの無駄な計算をしていることは想像に難くないです。実行速度が落ちている理由はこれです。
上記の非効率を解決するために重要になるのがメモ化です。
メモ化とは
上記のアルゴリズムで問題だったのは同じ計算を繰り返していたからでした。なので、一度計算したものはキャッシュしておいて、2回目以降はキャッシュした値を使うことにします。これがメモ化です。
概念的には下記のコールグラフのようになります。
最初にfibonacci(4) と fibonacci(3) を計算するのは同じですが、その結果をキャッシュしておき、2回目以降は計算を省くことで計算速度を速めます。コールグラフでは計算が省かれたものは薄く表示しています。
説明は以上にして、メモ化を実装していきましょう。
メモ化の実装(素朴な方法)
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)
ですが、それを見つけるのに時間がかかってしまいます。
しかも関数の外に連想配列を定義しておかなければなりません。
もう少し良い方法はないでしょうか。
メモ化の実装(スッキリ書ける方法)
ここからがキモです。
下記の様にわかりやすくて読みやすい関数を目指します。
let fibonacci = memoize { fibonacci, n in
n < 2 ? n : fibonacci(n - 1) + fibonacci(n - 2)
}
上記ではmemoizeという関数をつかっています。これはSwiftの標準ライブラリの関数ではありません。
このmemoizeは自前の関数であり、これこそが、わかりやすくて読みやすいメモ化を実現してくれます。
memorize part 1
関数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の使い方をみてみます。
まずは素数判定関数で試してみます。
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(_:)の結果をメモ化してくれました。
(素数判定アルゴリズムが非常に非効率なものですがこの記事の主題ではないのでツッコミはナシでお願いします!)
フィボナッチ数でも試してみましょう。
// 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で初期化しておきましょう。
// 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の第二弾です。これが完成形です。
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の内部に押し込んでいます。
使い方は以下です。
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は同じ関数を参照しています。
関数が参照型であることを上手く使っています。
その他の計算例
階乗
メモ化の応用例としてよく使われる階乗の計算の利用方法を試してみます。
let factorial = memoize { factorial, n in
n < 2 ? n : n * factorial(n - 1)
}
factorial(5) // 120
factorial(7) // 5040
素数判定
isPrime(_:)は再帰的な構造ではないので第二引数のみ使います。
let isPrimeMemoized = memoize { _, n in
return isPrime(n)
}
isPrimeMemoized(10) // false
参考にした資料
2014年のWWDCのAdvanced Swiftというセッションを参考にしました。
https://developer.apple.com/videos/play/wwdc2014/404/