moratorium

以下がコードです。

import Complex

fft_rec c len step sign =
  if len == 2
  then [((c !! 0) + (c !! step)), ((c !! 0) - (c !! step ))]
  else
    concat [(map gen_head arg_lst), (map gen_tail arg_lst)]
    where
      halflen = floor ((fromIntegral len) / (fromIntegral 2))
      fft0    = fft_rec (drop    0 c) halflen (step * 2) sign
      fft1    = fft_rec (drop step c) halflen (step * 2) sign
      arg_lst = zip3 (take halflen (iterate (+1) 0)) fft0 fft1
      gen_i i = 2.0 * pi * i / (fromIntegral len)
      gen_w i = (cos $ (gen_i i)) :+ (sign * (sin $ (gen_i i)))
      gen_head (i, f0, f1) = (f0 + (gen_w i) * f1)
      gen_tail (i, f0, f1) = (f0 - (gen_w i) * f1)

fft c =
  fft_rec c (length c) 1 1

ifft c =
  map divlen (fft_rec c (length c) 1 (-1.0))
  where
    divlen f = (((realPart f) / (fromIntegral (length c)) + 0.5)
                :+ ((imagPart f) / (fromIntegral (length c)) + 0.5))

convolution lst1 lst2 =
  map (\(c1, c2) -> c1 * c2) (zip lst1 lst2)

fft_mul c1 c2 =
  ifft (convolution (fft c1) (fft c2))

p [] = putStrLn ""
p (x:xs) =
  do putStr (show (floor (realPart x)))
     p xs

main = do
    p (reverse (fft_mul
        (reverse [(0.0 :+ 0.0), (0.0 :+ 0.0), (1.0 :+ 0.0), (2.0 :+ 0.0)])
        (reverse [(0.0 :+ 0.0), (0.0 :+ 0.0), (2.0 :+ 0.0), (3.0 :+ 0.0)])))

実行結果。

[kzk@dev]~% ghc fft.hs; ./a.out
0276

main部分では12 * 23を計算しています。fft_mulがかなり綺麗に書けているのが美しい。Ｃで書いた時にはメモリ周りのバグでも苦しんだので、その辺リストにして何も気にせず書けるのが良い。まぁ、遅いんだけど...。