1 files changed, 77 insertions, 0 deletions
diff --git a/2021/14.hs b/2021/14.hs
new file mode 100644
index 0000000..b088a8f
--- /dev/null
+++ b/2021/14.hs
@@ -0,0 +1,77 @@
+module Main where
+
+import qualified Data.Array as A
+import Data.List
+import qualified Data.Map.Strict as Map
+import Data.Map.Strict (Map)
+
+import Input
+import Util
+
+
+step :: Eq a => [((a, a), a)] -> [a] -> [a]
+step rules (a:b:s) = case lookup (a, b) rules of
+                       Nothing -> a : step rules (b:s)
+                       Just c -> a : c : step rules (b:s)
+step _ l = l
+
+maphistogram :: Ord a => [a] -> Map a Int
+maphistogram = foldr (\x -> Map.insertWith (+) x 1) mempty
+
+-- Dynamic programming: table
+--   (all letter pairs) x (number of steps) x (all letter pairs)
+-- where
+--   table[p][n][q] means the number of times that q appears in the string
+--   produced when stepping n times starting from the string p.
+-- Defining relation:
+--   table[p][0][q] = if p == q then 1 else 0
+--   table[p][n][q] = sum over pairs r:
+--                      table[p][n-1][r] * (number of occurrences of q in step(r))
+dyn :: [((Char, Char), Char)]
+    -> Int
+    -> (Map Char Int, Int, A.Array ((Int, Int), Int, (Int, Int)) Int)
+dyn rules nsteps = (charmap, maxid, table)
+  where
+    chars = uniq . sort $ concatMap (\((a,b),c) -> [a,b,c]) rules
+    charlist = zip chars [0::Int ..]
+    charmap = Map.fromList charlist
+    maxid = snd (fst (Map.deleteFindMax charmap))
+    table = A.array (((0, 0), 0, (0, 0)), ((maxid, maxid), nsteps, (maxid, maxid)))
+        [(((p1, p2), n, (q1, q2))
+         ,if n == 0 then if p == q then 1 else 0
+                    else sum [table A.! (p, n - 1, r) * count (q1c, q2c) (step rules [r1c, r2c])
+                             | (r1c, r1) <- charlist, (r2c, r2) <- charlist
+                             , let r = (r1, r2)])
+        | p1 <- [0..maxid], p2 <- [0..maxid]
+        , let p = (p1, p2)
+        , n <- [0..nsteps]
+        , (q1c, q1) <- charlist, (q2c, q2) <- charlist
+        , let q = (q1, q2)]
+    count (a, b) (x:y:s) | a == x, b == y = 1 + count (a, b) (y:s)
+                         | otherwise = count (a, b) (y:s)
+    count _ _ = 0
+
+solve :: [((Char, Char), Char)] -> String -> Int -> Int
+solve rules template nsteps =
+    let (charmap, maxid, table) = dyn rules nsteps
+        inttemplate = map (charmap Map.!) template
+        hist = Map.fromListWith (+) $ concat
+            [let n = table A.! ((a, b), nsteps, (c, d))
+             in [(c, n), (d, n)]
+            | (a, b) <- zip inttemplate (tail inttemplate)
+            , c <- [0..maxid]
+            , d <- [0..maxid]]
+        hist' = Map.map (`div` 2) . Map.unionWith (+) (Map.fromList [(charmap Map.! head template, 1), (charmap Map.! last template, 1)]) $ hist
+        -- freqs = Map.assocs hist
+    in -- trace (show table ++ "\n" ++ show freqs) $
+       maximum (Map.elems hist') - minimum (Map.elems hist')
+
+main :: IO ()
+main = do
+    template : _ : rules' <- getInput 14
+    let rules = map (\[a,b,' ','-','>',' ',c] -> ((a, b), c)) rules'
+    -- let string1 = iterate (step rules) template !! 10
+    --     hist1 = Map.elems (maphistogram string1)
+    -- print (maximum hist1 - minimum hist1)
+    print (solve rules template 10)
+    print (solve rules template 40)