{-# Language ImportQualifiedPost, QuasiQuotes #-}
{-|
Module      : Main
Description : Day 14 solution
Copyright   : (c) Eric Mertens, 2021
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2021/day/14>

Build a huge polymer chain and compute how many of
each element it contains.

This problem requires memoization as the size of the
resulting polymer would be humongous!

-}
module Main (main) where

import Advent (format, power, counts)
import Data.Map (Map)
import Data.Map.Strict qualified as Map

-- | >>> :main
-- 2068
-- 2158894777814
main :: IO ()
IO ()
main =
 do ([Char]
seed, [(Char, Char, Char)]
table) <- [format|2021 14 %s%n%n(%c%c -> %c%n)*|]
    let rule :: Map (Char, Char) (Map (Char, Char) Int)
rule = [(Char, Char, Char)] -> Map (Char, Char) (Map (Char, Char) Int)
forall a. Ord a => [(a, a, a)] -> Map (a, a) (Map (a, a) Int)
tableToRule [(Char, Char, Char)]
table
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Map (Char, Char) (Map (Char, Char) Int) -> Integer -> [Char] -> Int
forall a.
Ord a =>
Map (a, a) (Map (a, a) Int) -> Integer -> [a] -> Int
solve Map (Char, Char) (Map (Char, Char) Int)
rule Integer
10 [Char]
seed)
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Map (Char, Char) (Map (Char, Char) Int) -> Integer -> [Char] -> Int
forall a.
Ord a =>
Map (a, a) (Map (a, a) Int) -> Integer -> [a] -> Int
solve Map (Char, Char) (Map (Char, Char) Int)
rule Integer
40 [Char]
seed)

solve :: Ord a => Map (a,a) (Map (a,a) Int) -> Integer -> [a] -> Int
solve :: forall a.
Ord a =>
Map (a, a) (Map (a, a) Int) -> Integer -> [a] -> Int
solve Map (a, a) (Map (a, a) Int)
rule Integer
n [a]
seed = Map a Int -> Int
forall a. Ord a => Map a a -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum Map a Int
occ Int -> Int -> Int
forall a. Num a => a -> a -> a
- Map a Int -> Int
forall a. Ord a => Map a a -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum Map a Int
occ
  where
    ruleN :: Map (a, a) (Map (a, a) Int)
ruleN = (Map (a, a) (Map (a, a) Int)
 -> Map (a, a) (Map (a, a) Int) -> Map (a, a) (Map (a, a) Int))
-> Map (a, a) (Map (a, a) Int)
-> Integer
-> Map (a, a) (Map (a, a) Int)
forall a. HasCallStack => (a -> a -> a) -> a -> Integer -> a
power ((Map (a, a) Int -> Map (a, a) Int)
-> Map (a, a) (Map (a, a) Int) -> Map (a, a) (Map (a, a) Int)
forall a b. (a -> b) -> Map (a, a) a -> Map (a, a) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Map (a, a) Int -> Map (a, a) Int)
 -> Map (a, a) (Map (a, a) Int) -> Map (a, a) (Map (a, a) Int))
-> (Map (a, a) (Map (a, a) Int)
    -> Map (a, a) Int -> Map (a, a) Int)
-> Map (a, a) (Map (a, a) Int)
-> Map (a, a) (Map (a, a) Int)
-> Map (a, a) (Map (a, a) Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map (a, a) (Map (a, a) Int) -> Map (a, a) Int -> Map (a, a) Int
forall a b.
(Ord a, Ord b) =>
Map a (Map b Int) -> Map a Int -> Map b Int
applyRule) Map (a, a) (Map (a, a) Int)
rule Integer
n

    start :: Map (a, a) Int
start = [(a, a)] -> Map (a, a) Int
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Map a Int
counts ([a] -> [a] -> [(a, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
seed ([a] -> [a]
forall a. HasCallStack => [a] -> [a]
tail [a]
seed))

    occ :: Map a Int
occ = (Int -> Int -> Int) -> a -> Int -> Map a Int -> Map a Int
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ([a] -> a
forall a. HasCallStack => [a] -> a
head [a]
seed) Int
1
        (Map a Int -> Map a Int) -> Map a Int -> Map a Int
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> ((a, a) -> a) -> Map (a, a) Int -> Map a Int
forall k2 a k1.
Ord k2 =>
(a -> a -> a) -> (k1 -> k2) -> Map k1 a -> Map k2 a
Map.mapKeysWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) (a, a) -> a
forall a b. (a, b) -> b
snd
        (Map (a, a) Int -> Map a Int) -> Map (a, a) Int -> Map a Int
forall a b. (a -> b) -> a -> b
$ Map (a, a) (Map (a, a) Int) -> Map (a, a) Int -> Map (a, a) Int
forall a b.
(Ord a, Ord b) =>
Map a (Map b Int) -> Map a Int -> Map b Int
applyRule Map (a, a) (Map (a, a) Int)
ruleN Map (a, a) Int
start

-- | Generate a replacement rule map from a list of input productions
--
-- >>> tableToRule [('L','R','M')] -- LR -> M
-- fromList [(('L','R'),fromList [(('L','M'),1),(('M','R'),1)])]
tableToRule :: Ord a => [(a,a,a)] -> Map (a,a) (Map (a,a) Int)
tableToRule :: forall a. Ord a => [(a, a, a)] -> Map (a, a) (Map (a, a) Int)
tableToRule [(a, a, a)]
xs = [((a, a), Map (a, a) Int)] -> Map (a, a) (Map (a, a) Int)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [((a
l,a
r), [(a, a)] -> Map (a, a) Int
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Map a Int
counts [(a
l,a
m), (a
m,a
r)]) | (a
l,a
r,a
m) <- [(a, a, a)]
xs]

-- | Apply a replacement rule to a map of counts.
--
-- >>> :set -XOverloadedLists
-- >>> applyRule [('a', [('b',1),('c',2)]),('z',[('y',1)])] [('a',10)]
-- fromList [('b',10),('c',20)]
applyRule :: (Ord a, Ord b) => Map a (Map b Int) -> Map a Int -> Map b Int
applyRule :: forall a b.
(Ord a, Ord b) =>
Map a (Map b Int) -> Map a Int -> Map b Int
applyRule Map a (Map b Int)
r Map a Int
m = (Int -> Int -> Int) -> [Map b Int] -> Map b Int
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [(Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
*) (Int -> Int) -> Map b Int -> Map b Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map a (Map b Int)
r Map a (Map b Int) -> a -> Map b Int
forall k a. Ord k => Map k a -> k -> a
Map.! a
k | (a
k,Int
v) <- Map a Int -> [(a, Int)]
forall k a. Map k a -> [(k, a)]
Map.toList Map a Int
m]