Haskellのメモ化について

GCJのQualification Round出ました。
最近Haskellをはじめたので、Haskellで参戦しました。
結果とか感想は結果がでてからとして、本番中に困ったことが起こりました。
それは・・・

Haskellでメモ化ってどうするの・・・?

ググッてもモナドなるものを使う方法しかでません。
Haskell入門書読んでもまだいまいち理解出来ていないモナドを使うのはちょっと。
ということで本番中にあれこれ悩みつつ、思い浮かんだ方法をメモろうと思います。
基本的に、考えていた思考どおりの流れで書いていきます。
(※上記のとおり初心者なのです。Haskellerとして正しいかは知りませんし、逆にもしかしたら超常識なのかも知れません)

実際にGCJで使ったコードは汚いし、まだ開催中なので他のコードで書きます。
簡単なところでフィボナッチ関数(もちろん再帰的な定義で)を考えてみます。

まずは普通に書いてみましょう。

fib :: Int -> Int
fib 0 = 1
fib 1 = 1
fib n = fib (n - 1) + fib (n - 2)

まあ普通ですね。fib 40を計算してみると、手元の環境では10秒ほどかかりました。

さてメモ化です・・・。
まず重要なのは、Haskellでは変数への代入ができません。
ということは普段C++で使っているような。

int memo[100]; // initialized with 0
int fib(int n){
  if(n < 2) return 1;
  else if(memo[n] == 0) return memo[n] = fib(n-1) + fib(n-2);
  else memo[n];
}

みたいな戦法は使えません。さあ、どうする・・・

まあHaskellだし、どうせリストやろー

ということでリストにメモっていく。
とりあえずコード↓

memofib :: [Int]
memofib = memofib' 0
    where
      memofib' n = let val = if n < 2 then 1 else memofib !! (n-1) + memofib !! (n-2)
                   in [val] ++ memofib' (n+1)

fib :: Int -> Int
fib n = memofib !! n

はいきたー。こんどはfib 40が一瞬で答えが出ます。
memofibは無限リストで、そこに答えをメモっていきます。前に計算した値を使う(実際には、遅延評価のため後に計算される?下記参照)ときは(!!)演算子を使います。
(本番中は、これはfib 0、fib1...を順に計算していてDPのような構造になっているのではと思いましたが、Haskellの遅延評価のおかげでこれでメモ化のようになってい・・・ますよね、多分・・・。つまりリストの必要ない部分は計算していないはずです。)
はいっ、終わり!としたいところですが、使っていた場面が結構性能にクリティカルなところで、リストだとランダムアクセスはできないので時間的に微妙な気がしました。

じゃあ配列やろー

ということで配列を使ってみましょう。
Haskellでは、arrayという関数(?)をつかって配列を作るみたいです(知らなかった;;)
初期化には、インデックスとのタプルを使うので、上記の無限リストを使って初期化できそうです。
ということで書いてみる。

import Data.Array (Array,array,(!))

memofibArray :: Array Int Int
memofibArray = array (0, 100) $ zip [0..100] memofib
    where
      memofib :: [Int]
      memofib = memofib' 0
      memofib' n = let val = if n < 2 then 1 else memofibArray ! (n-1) + memofibArray ! (n-2)
                   in [val] ++ memofib' (n+1)

fib :: Int -> Int
fib n = memofibArray ! n

こんな感じで書けました。memofibArrayで包んで、memofib'でmemofibを読んでたところをmemofibArrayから読むようにしただけです。
Arrayからの読み出しは(!)演算子を使います。
この例だとパフォーマンスの差はほとんどないですが、もっと複雑な例になると効いてくるんではないでしょうか。

さて、これでメモ化は基本的には完成ですね。
実際に、GCJ中にはここで終わりました。
でも・・・

最初の定義から結構離れてるし、無限リストとかムズイ><

つまり直感的な定義とはかけ離れた定義になり、とっさに書きづらい。・・・と個人的には思います。
そこで、終わったあとに次のような定義を思い浮かびました。

import Data.Array (Array,array,(!))

fib :: Int -> Int
fib = let memofibArray = array (0, 100) [(i, fib' i) | i <- [0..100]]
        in (!) memofibArray
         where
           fib' :: Int -> Int
           fib' 0 = 1
           fib' 1 = 1
           fib' n = fib (n - 1) + fib (n - 2)

中に一番最初の定義と同じ形がでてきているのが分かります。
無限リストも出てきていません。
(最初の定義の形を使うのはもっと最初の方から出来ましたが・・・w)
これで直感的にメモ化が出来るのではないでしょうか。

ちなみに↓だと上手く動かないようでよく分かりません。
これだとmemofibArrayが毎回生成されるんですかね。

import Data.Array (Array,array,(!))

-- This code doesn't work well
fib :: Int -> Int
fib n = let memofibArray = array (0, 100) [(i, fib' i) | i <- [0..100]]
        in memofibArray ! n
         where
           fib' :: Int -> Int
           fib' 0 = 1
           fib' 1 = 1
           fib' n = fib (n - 1) + fib (n - 2)

(追記)
関数型言語らしく抽象化した

import Data.Array (Array, Ix, array, (!), range)

memorize :: Ix a => (a -> b) -> (a, a) -> (a -> b)
memorize func rng = let memoArray = array rng [(i, func i) | i <- (range rng)]
                    in (!) memoArray

fib :: Int -> Int
fib = memorize fib' (0, 100)
    where
      fib' :: Int -> Int
      fib' 0 = 1
      fib' 1 = 1
      fib' n = fib (n - 1) + fib (n - 2)