122 lines
4.2 KiB
Haskell
122 lines
4.2 KiB
Haskell
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)
|
|
|
|
main :: IO ()
|
|
main = do
|
|
input <- getContents
|
|
let
|
|
polymer = parsePolymer input
|
|
(putStrLn . show . solution1) polymer
|
|
(putStrLn . show . solution2) polymer
|
|
|
|
parsePolymer :: String -> ([Char], Map [Char] Char)
|
|
parsePolymer input = (template input, rules input)
|
|
where
|
|
template :: String -> String
|
|
template = head . splitByString "\n\n"
|
|
rules :: String -> Map [Char] Char
|
|
rules = Map.fromList . map tupelize . map (splitByString " -> ") . lines . last . splitByString "\n\n"
|
|
tupelize :: [String] -> ([Char], Char)
|
|
tupelize untupled = (head untupled, (head . last) untupled)
|
|
|
|
solution1 :: ([Char], Map [Char] Char) -> Int
|
|
solution1 (template, rules) = uncurry (-) (leastMostOcc (polymerizeN rules template 10))
|
|
|
|
solution2 :: ([Char], Map [Char] Char) -> Int
|
|
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
|
|
where
|
|
step :: [Char] -> [Char]
|
|
step [] = []
|
|
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 [
|
|
"NNCB",
|
|
"",
|
|
"CH -> B",
|
|
"HH -> N",
|
|
"CB -> H",
|
|
"NH -> C",
|
|
"HB -> C",
|
|
"HC -> B",
|
|
"HN -> C",
|
|
"NN -> C",
|
|
"BH -> H",
|
|
"NC -> B",
|
|
"NB -> B",
|
|
"BN -> B",
|
|
"BB -> N",
|
|
"BC -> B",
|
|
"CC -> N",
|
|
"CN -> C"]
|
|
|
|
parsedTestInput1 = parsePolymer testInput1
|
|
test1 = solution1 parsedTestInput1 -- 1588
|
|
test2 = solution2 parsedTestInput1 -- 2188189693529
|