Polynomial stuff

app/Main.hs

@@ -1,6 +1,13 @@
-module Main where
+module Main (main) where
+
+import Data.Vector qualified as V
+import Poly
 
 main :: IO ()
 main = do
-  putStrLn "Hello, Haskell!"
-
+  let f :: Poly Int = Poly (V.fromList [1, 2, 3])
+  let g :: Poly Int = Poly (V.fromList [4, 5])
+  putStrLn $ "f: " <> show f <> ", g: " <> show g
+  putStrLn $ "f + g: " <> show (f + g)
+  putStrLn $ "f * g: " <> show (f * g)
+  putStrLn $ "Karatsuba f * g: " <> show (normalize $ karatsubaMult f g)

mathematical-algorithms.cabal

@@ -16,11 +16,23 @@ build-type:         Simple
 common warnings
     ghc-options: -Wall
 
+common deps
+    build-depends:
+        base ^>=4.17.2.1,
+        vector,
+
+library
+    import:           warnings, deps
+    exposed-modules:
+        Poly
+    hs-source-dirs:   src
+    default-language: GHC2021
+
 executable mathematical-algorithms
-    import:           warnings
+    import:           warnings, deps
     main-is:          Main.hs
     -- other-modules:
-    -- other-extensions:
-    build-depends:    base ^>=4.17.2.1
+    build-depends:
+        mathematical-algorithms
     hs-source-dirs:   app
     default-language: GHC2021

src/Poly.hs

@@ -0,0 +1,93 @@
+module Poly where
+
+import Data.List
+import Data.Maybe
+import Data.Vector qualified as V
+
+-- Zip two vectors while padding 0s on the shorter vector.
+vecZipPad0With :: (Num a) => (a -> a -> a) -> V.Vector a -> V.Vector a -> V.Vector a
+vecZipPad0With f xs ys = V.generate (max (V.length xs) (V.length ys)) $
+  \i -> fromMaybe 0 (xs V.!? i) `f` fromMaybe 0 (ys V.!? i)
+
+-- | Polynomial type.
+--
+-- >>> Poly (V.fromList [1 .. 5])
+-- 1 X^0 + 2 X^1 + 3 X^2 + 4 X^3 + 5 X^4
+
+-- >>> Poly (V.fromList [1, 2]) * Poly (V.fromList [3, 4, 5])
+-- 3 X^0 + 10 X^1 + 13 X^2 + 10 X^3
+
+-- >>> Poly (V.fromList [1, 2]) * Poly (V.fromList [])
+-- 0 X^0
+newtype Poly a = Poly (V.Vector a)
+  deriving (Eq)
+
+-- | Degree, assuming top term is nonzero
+degree :: Poly a -> Int
+degree (Poly f) = length f - 1
+
+-- | Shift up polynomial by X^n
+shiftUp :: (Num a) => Int -> Poly a -> Poly a
+shiftUp n (Poly f) = Poly $ V.replicate n 0 <> f
+
+-- | Shift down polynomial by X^n
+shiftDown :: Int -> Poly a -> Poly a
+shiftDown n (Poly f) = Poly $ V.drop n f
+
+-- | Remainder under X^n
+remXn :: Int -> Poly a -> Poly a
+remXn n (Poly f) = Poly $ V.take n f
+
+-- | Normalize polynomial, removing leading 0s
+-- 
+-- >>> normalize $ Poly (V.fromList [1, 0, 0])
+-- 1 X^0
+-- 
+-- >>> normalize $ Poly (V.fromList [1, 2, 3, 0])
+-- 1 X^0 + 2 X^1 + 3 X^2
+normalize :: (Eq a, Num a) => Poly a -> Poly a
+normalize (Poly f) = Poly remain
+  where
+    (_, remain) = V.spanR (== 0) f
+
+-- | This Num instance implements the classical multiplication.
+instance (Num a) => Num (Poly a) where
+  (+) :: Poly a -> Poly a -> Poly a
+  Poly f + Poly g = Poly $ vecZipPad0With (+) f g
+  (-) :: Poly a -> Poly a -> Poly a
+  Poly f - Poly g = Poly $ vecZipPad0With (-) f g
+  (*) :: Poly a -> Poly a -> Poly a
+  Poly f * Poly g = sum (Poly <$> mults)
+    where
+      mults = V.imap (\i fi -> V.map (fi *) (V.replicate i 0 <> g)) f
+  negate :: Poly a -> Poly a
+  negate (Poly f) = Poly $ V.map negate f
+  abs :: Poly a -> Poly a
+  abs = error "abs: invalid on poly"
+  signum :: Poly a -> Poly a
+  signum = error "signum: invalid on poly"
+  fromInteger :: Integer -> Poly a
+  fromInteger = Poly . V.singleton . fromInteger
+
+instance (Show a) => Show (Poly a) where
+  show :: (Show a) => Poly a -> String
+  show (Poly p) = intercalate " + " . V.toList $ V.imap (\i coeff -> show coeff <> " X^" <> show i) p
+
+karatsubaMult :: (Num a) => Poly a -> Poly a -> Poly a
+karatsubaMult a b = atLog degBound a b
+  where
+    degBound = fromJust $ find (> max (degree a) (degree b)) [2 ^ i | i <- [0 :: Int ..]]
+
+    -- degBnd: power-of-two degree bound
+    atLog degBnd f g = case degBnd of
+      1 -> f * g
+      _ -> shiftUp degBnd prod1 + shiftUp nextBound (prodAdd - prod0 - prod1) + prod0
+      where
+        nextBound = degBnd `div` 2
+        f1 = shiftDown nextBound f
+        f0 = remXn nextBound f
+        g1 = shiftDown nextBound g
+        g0 = remXn nextBound g
+        prod0 = atLog nextBound f0 g0
+        prod1 = atLog nextBound f1 g1
+        prodAdd = atLog nextBound (f0 + f1) (g0 + g1)