ATC001-C: Implement NTT

Author: minoki

atc001-c/Main.hs

@@ -37,19 +37,22 @@ main = do
 -- Fast Fourier Transform (FFT)
 --
 
+halve :: G.Vector vec a => vec a -> vec a
+halve v = let n = G.length v
+          in G.generate (n `quot` 2) $ \j -> v G.! (j * 2)
+
 fft :: forall vec a. (Num a, G.Vector vec a)
-    => vec a -- ^ For a primitive n-th root of unity @u@, @[1,u,u^2 .. u^(n-1)]@
+    => [vec a] -- ^ For a primitive n-th root of unity @u@, @iterate halve [1,u,u^2 .. u^(n-1)]@
     -> vec a -- ^ a polynomial of length n (= 2^k for some k)
     -> vec a
-fft u f | n == 1 = f
-        | otherwise = let !n2 = n `quot` 2
-                          r0, r1', u2, t0, t1' :: vec a
-                          r0 = G.generate n2 $ \j -> (f G.! j) + (f G.! (j + n2))
-                          r1' = G.generate n2 $ \j -> ((f G.! j) - (f G.! (j + n2))) * u G.! j
-                          !u2 = G.generate n2 $ \j -> u G.! (j * 2)
-                          !t0 = fft u2 r0
-                          !t1' = fft u2 r1'
-                      in G.generate n $ \j -> if even j then t0 G.! (j `quot` 2) else t1' G.! (j `quot` 2)
+fft (u:u2) f | n == 1 = f
+             | otherwise = let !n2 = n `quot` 2
+                               r0, r1', t0, t1' :: vec a
+                               r0 = G.generate n2 $ \j -> (f G.! j) + (f G.! (j + n2))
+                               r1' = G.generate n2 $ \j -> ((f G.! j) - (f G.! (j + n2))) * u G.! j
+                               !t0 = fft u2 r0
+                               !t1' = fft u2 r1'
+                           in G.generate n $ \j -> if even j then t0 G.! (j `quot` 2) else t1' G.! (j `quot` 2)
   where n = G.length f
 
 mulFFT :: U.Vector Int -> U.Vector Int -> U.Vector Int
@@ -59,6 +62,7 @@ mulFFT !f !g = let n' = U.length f + U.length g - 2
                    n = bit k
                    u :: U.Vector (Complex Double)
                    u = U.generate n $ \j -> cis (fromIntegral j * (2 * pi / fromIntegral n))
+                   us = iterate halve u
                    f' = U.generate n $ \j -> if j < U.length f then
                                                fromIntegral (f U.! j)
                                              else
@@ -67,10 +71,10 @@ mulFFT !f !g = let n' = U.length f + U.length g - 2
                                                fromIntegral (g U.! j)
                                              else
                                                0
-                   f'' = fft u f'
-                   g'' = fft u g'
+                   f'' = fft us f'
+                   g'' = fft us g'
                    fg = U.generate n $ \j -> (f'' U.! j) * (g'' U.! j)
-                   fg' = fft (U.map conjugate u) fg
+                   fg' = fft (map (U.map conjugate) us) fg
                in U.generate n $ \j -> round (realPart (fg' U.! j) / fromIntegral n)
 
 --