day14: Complete second puzzle

14/Itertools.hs

@@ -0,0 +1,14 @@
+module Itertools (
+  flatmap,
+  windows
+  ) where
+
+windows :: Int -> [a] -> [[a]]
+windows _ [] = []
+windows i (x:xs)
+  | length xs < i-1 = []
+  | otherwise = (x:take (i-1) xs):windows i xs
+
+flatmap :: (t -> [a]) -> [t] -> [a]
+flatmap _ []     = []
+flatmap f (x:xs) = f x ++ flatmap f xs

14/polymerization.hs

@@ -1,5 +1,8 @@
+import           Data.List       (maximumBy, minimumBy)
 import           Data.Map.Strict (Map, (!))
 import qualified Data.Map.Strict as Map
+import           Data.Ord        (comparing)
+import           Itertools       (flatmap, windows)
 import           Parsing         (splitByString)
 import           Reducers        (leastMostOcc)
 
@@ -9,6 +12,7 @@ main = do
   let
     polymer = parsePolymer input
   (putStrLn . show . solution1) polymer
+  (putStrLn . show . solution2) polymer
 
 parsePolymer :: String -> ([Char], Map [Char] Char)
 parsePolymer input = (template input, rules input)
@@ -21,13 +25,18 @@ parsePolymer input = (template input, rules input)
     tupelize untupled = (head untupled, (head . last) untupled)
 
 solution1 :: ([Char], Map [Char] Char) -> Int
-solution1 (template, rules) = uncurry (-) (leastMostOcc polymer)
-  where
-    polymer :: [Char]
-    polymer = iterate (polymerize rules) template !! 10
+solution1 (template, rules) = uncurry (-) (leastMostOcc (polymerizeN rules template 10))
 
 solution2 :: ([Char], Map [Char] Char) -> Int
-solution2 (template, rules) = undefined
+solution2 (template, rules) = ((snd . most) polyCounts) - ((snd . least) polyCounts)
+  where
+    polyCounts :: Map Char Int
+    polyCounts = polyCounter rules template 40
+    least = minimumBy (comparing snd) . Map.assocs
+    most = maximumBy (comparing snd) . Map.assocs
+
+polymerizeN :: Map [Char] Char -> [Char] -> Int -> [Char]
+polymerizeN rules template i = iterate (polymerize rules) template !! i
 
 polymerize :: Map [Char] Char -> [Char] -> [Char]
 polymerize rules template = step template
@@ -37,6 +46,54 @@ polymerize rules template = step template
     step [final] = [final]
     step (x:xs)  = x:(rules ! [x, head xs]):(step xs)
 
+-- Since we don't actually need to synthesize the polymer, but only
+-- need to know how many of each elements occur, we can get around
+-- having to create a ridiculously large list by operating on the rule
+-- pairs instead.
+--
+-- E.g. NNCB -> NCNBCHB (see testInput1), in this case we turn:
+--
+--   - NN -> NC, CN
+--   - NC -> NB, BC
+--   - CB -> CH, HB
+--
+-- All we need to do is keep track of the counts of these
+-- sub-sequences, not of the full polymer.
+--
+-- To get back to the individual elements, we then just need to
+-- de-window it all again.
+polyCounter :: Map [Char] Char -> [Char] -> Int -> Map Char Int
+polyCounter rules template iterations = elCounts
+  where
+    initial :: Map [Char] Int
+    initial = (Map.fromListWith (+) . map (\e -> (e, 1)) . windows 2) template
+    step :: Map [Char] Int -> Map [Char] Int
+    step = Map.fromListWith (+) . flatmap resultingPolies . Map.assocs
+    resultingPolies :: ([Char], Int) -> [([Char], Int)]
+    resultingPolies (p, i) = [([head p, c], i), ([c, last p], i)]
+      where c = rules ! p
+    polyCounts :: Map [Char] Int
+    polyCounts = iterate step initial !! iterations
+
+    -- Because the windows are size 2, every element of a sub-binding is
+    -- counted twice (e.g. NBCBC -> NB, BC, CB, BC), except the first and
+    -- last element. We know the first and last elements (because we
+    -- operated on windows, they'll match the first and last of the
+    -- template), so we can un-window the whole batch by dividing their
+    -- number of occurrences by two, except for the first and last, where
+    -- we subtract the number by 1 first and add 1 back later.
+    elCounts :: Map Char Int
+    elCounts = Map.mapWithKey deWindow naiveSum
+      where
+        els :: ([Char], Int) -> [(Char, Int)]
+        els (p, i) = [(head p, i), (last p, i)]
+        naiveSum :: Map Char Int
+        naiveSum = (Map.fromListWith (+) . flatmap els . Map.assocs) polyCounts
+        deWindow :: Char -> Int -> Int
+        deWindow c i
+          | c == head template || c == last template = (i - 1) `quot` 2 + 1
+          | otherwise = i `quot` 2
+
 -- Tests
 
 testInput1 = unlines [
@@ -59,6 +116,6 @@ testInput1 = unlines [
   "CC -> N",
   "CN -> C"]
 
-testInputParsed1 = parsePolymer testInput1
-test1 = solution1 testInputParsed1 -- 1588
-test2 = solution2 testInputParsed1 -- 2188189693529
+parsedTestInput1 = parsePolymer testInput1
+test1 = solution1 parsedTestInput1 -- 1588
+test2 = solution2 parsedTestInput1 -- 2188189693529