advent14/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2021/12/16/advent-of-code-2021-day-14/
   2
   3 import Data.Text ()
   4 import qualified Data.Text.IO as TIO
   5
   6 import Data.Attoparsec.Text
   7 import Control.Applicative
   8
   9 import Data.List
  10 import qualified Data.Map as M
  11 import Data.Map ((!))
  12 import qualified Data.MultiSet as MS
  13 import qualified Data.Set as S
  14
  15 type RuleSet = M.Map String String
  16 type PolyPairs = MS.MultiSet String
  17
  18
  19 main :: IO ()
  20 main =
  21   do  text <- TIO.readFile "data/advent14.txt"
  22       let (template, rules) = successfulParse text
  23       print $ part1 rules template
  24       print $ part2 rules template
  25
  26 part1 :: RuleSet -> String -> Int
  27 part1 rules template = (last counts) - (head counts)
  28   where result = (simulateNaive rules template) !! 10
  29         counts = sort $ map snd $ MS.toOccurList $ MS.fromList result
  30
  31 simulateNaive :: RuleSet -> String -> [String]
  32 simulateNaive rules polymer = iterate (stepNaive rules) polymer
  33
  34 stepNaive :: RuleSet -> String -> String
  35 stepNaive rules polymer = merge polymer $ concatMap (rules !) $ mkPairs polymer
  36
  37 part2 :: RuleSet -> String -> Int
  38 part2 rules template = (last counts) - (head counts)
  39   where pairs = MS.fromList $ mkPairs template
  40         result = (simulate rules pairs) !! 40
  41         elementCounts = countElements result
  42         -- counts = sort $ map snd $ MS.toOccurList elementCounts
  43         counts = sort $ M.elems elementCounts
  44
  45 simulate :: RuleSet -> PolyPairs -> [PolyPairs]
  46 simulate rules polymer = iterate (step rules) polymer
  47
  48 step :: RuleSet -> PolyPairs -> PolyPairs
  49 step rules polymer = MS.union firsts seconds
  50   where firsts = MS.map (addFirst rules) polymer
  51         seconds = MS.map (addSecond rules) polymer
  52
  53 addFirst :: RuleSet -> String -> String
  54 addFirst rules pair = a : c
  55   where a = pair!!0
  56         c = rules ! pair
  57
  58 addSecond :: RuleSet -> String -> String
  59 addSecond rules pair = c ++ [a]
  60   where a = pair!!1
  61         c = rules ! pair
  62
  63 countElements :: PolyPairs -> M.Map Char Int
  64 countElements pairs = counts
  65   where firsts = MS.map (!!0) pairs
  66         seconds = MS.map (!!1) pairs
  67         elems = S.union (MS.toSet firsts) (MS.toSet seconds)
  68         counts = M.map ((`div` 2) . (+ 1)) $ MS.toMap $ MS.union firsts seconds
  69
  70
  71 mkPairs :: String -> [String]
  72 mkPairs polymer = map stringify $ zip polymer $ tail polymer
  73 stringify (a, b) = [a, b]
  74
  75 merge :: [a] -> [a] -> [a]
  76 merge [] ys = ys
  77 merge (x:xs) ys = x : (merge ys xs)
  78
  79 -- Parse the input file
  80
  81 inputP = (,) <$> (many1 letter) <* many1 endOfLine <*> rulesP
  82
  83 rulesP = M.fromList <$> ruleP `sepBy` endOfLine
  84 ruleP = (,) <$> many1 letter <* " -> " <*> many1 letter
  85
  86 -- successfulParse :: Text -> (Integer, [Maybe Integer])
  87 successfulParse input =
  88   case parseOnly inputP input of
  89     Left  _err -> ("", M.empty) -- TIO.putStr $ T.pack $ parseErrorPretty err
  90     Right indata -> indata