Project Euler 511 残骸

p511.hs

import Data.Vector (Vector)
import qualified Data.Vector as V

newtype Zp = Z Int deriving (Show, Eq)

modulo :: Int
modulo = 10 ^ 9

z :: Int -> Zp
z a | a >= 0    = Z a'
    | otherwise = Z (modulo - a')
    where a' = a `rem` modulo

instance Num Zp where
    abs         = undefined
    signum      = undefined

    fromInteger = z . fromInteger

    negate (Z a)  = z (-a)
    (Z a) + (Z b) = z (a + b)
    (Z a) * (Z b) = z (a - b)

newtype Poly a = Poly { unPoly :: Vector a } deriving (Show, Eq)

instance Num a => Num (Poly a) where
    abs         = undefined
    signum      = undefined
    fromInteger = undefined
    negate      = Poly . V.map negate . unPoly
    a + b       = Poly $ V.zipWith (+) (unPoly a') (unPoly b')
                  where (a', b') = align a b
    a * b       = multiply a' b'
                  where (a', b') = align a b

align :: (Num a) => Poly a -> Poly a -> (Poly a, Poly a)
align a b
    | offset < 0 = (Poly . (V.++ padding) $ a', b)
    | otherwise  = (a, Poly . (V.++ padding) $ b')
    where a'      = unPoly a
          b'      = unPoly b
          offset  = V.length a' - V.length b'
          padding = V.replicate (abs offset) 0

degree :: Poly a -> Int
degree = V.length . unPoly

lift :: (Num a) => Int -> Poly a -> Poly a
lift n p = Poly . ((V.replicate n 0) V.++) . unPoly $ p

multiply :: (Num a) => Poly a -> Poly a -> Poly a
multiply a b
    | n' == 1   = Poly $ V.zipWith (*) (unPoly a) (unPoly b)
    | otherwise = p00 + (lift n p01) + (lift (n * 2) p11)
    where n'       = degree a
          n        = n' `quot` 2
          divide p = let (a, b) = V.splitAt n . unPoly $ p
                     in  (Poly a, Poly b)
          (a0, a1) = divide a
          (b0, b1) = divide b
          p00      = a0 `multiply` b0
          p11      = a1 `multiply` b1
          p01      = ((a0 + a1) `multiply` (b0 + b1)) - p00 - p11

main :: IO ()
main = print $ (p [1..4000]) * (p [1..4000])
       where p = (Poly . V.fromList) :: [Int] -> Poly Int