Data commit

Task/Knuths-power-tree/Haskell/knuths-power-tree.hs

@@ -0,0 +1,58 @@
+{-# LANGUAGE ScopedTypeVariables #-}
+
+module Rosetta.PowerTree
+    ( Natural
+    , powerTree
+    , power
+    )  where
+
+import           Data.Foldable (toList)
+import           Data.Map.Strict (Map)
+import qualified Data.Map.Strict as Map
+import           Data.Maybe (fromMaybe)
+import           Data.List (foldl')
+import           Data.Sequence (Seq (..), (|>))
+import qualified Data.Sequence as Seq
+import           Numeric.Natural (Natural)
+
+type M = Map Natural S
+type S = Seq Natural
+
+levels :: [M]
+levels = let s = Seq.singleton 1 in fst <$> iterate step (Map.singleton 1 s, s)
+
+step :: (M, S) -> (M, S)
+step (m, xs) = foldl' f (m, Empty) xs
+  where
+    f :: (M, S) -> Natural -> (M, S)
+    f (m', ys) n = foldl' g (m', ys) ns
+      where
+        ns :: S
+        ns = m' Map.! n
+
+        g :: (M, S) -> Natural -> (M, S)
+        g (m'', zs) k =
+            let l = n + k
+            in  case Map.lookup l m'' of
+                    Nothing -> (Map.insert l (ns |>  l) m'', zs |> l)
+                    Just _  -> (m'', zs)
+
+powerTree :: Natural -> [Natural]
+powerTree n
+    | n <= 0    = []
+    | otherwise = go levels
+  where
+    go :: [M] -> [Natural]
+    go []       = error "impossible branch"
+    go (m : ms) = fromMaybe (go ms) $ toList <$> Map.lookup n m
+
+power :: forall a. Num a => a -> Natural -> a
+power _ 0 = 1
+power a n = go a 1 (Map.singleton 1 a) $ tail $ powerTree n
+  where
+    go :: a -> Natural -> Map Natural a -> [Natural] -> a
+    go b _ _ []       = b
+    go b k m (l : ls) =
+        let b' = b * m Map.! (l - k)
+            m' = Map.insert l b' m
+        in  go b' l m' ls