advent12/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2023/12/15/advent-of-code-2023-day-12/
   2
   3 import AoC
   4 import Data.Text (Text)
   5 import qualified Data.Text.IO as TIO
   6 import Data.Attoparsec.Text hiding (take, takeWhile)
   7 import Control.Applicative
   8 import Data.List
   9 import qualified Data.Map.Strict as M
  10
  11 data Spring = Unknown | Damaged | Operational deriving (Show, Eq, Ord)
  12 data Record = Record [Spring] [Int] deriving (Show, Eq, Ord)
  13
  14 type Cache = M.Map Record Int
  15
  16
  17 main :: IO ()
  18 main =
  19   do  dataFileName <- getDataFileName
  20       text <- TIO.readFile dataFileName
  21       let records = successfulParse text
  22       -- print records
  23       print $ part1 records
  24       print $ part2 records
  25
  26 part1, part2 :: [Record] -> Int
  27 part1 = sum . fmap countViable
  28 part2 = sum . fmap (countViable . unfoldRecord)
  29
  30 unfoldRecord :: Record -> Record
  31 unfoldRecord (Record springs signature) = Record uSprings uSignature
  32   where uSprings = intercalate [Unknown] $ replicate 5 springs
  33         uSignature = concat $ replicate 5 signature
  34
  35 countViable :: Record -> Int
  36 countViable record = table M.! record
  37   where table0 = initialCache record
  38         table = fillTable table0 record
  39
  40 initialCache :: Record -> Cache
  41 initialCache (Record springs signature) = M.union lastOperational cache0
  42   where cache0 = M.union sprs sigs
  43         sprs = M.fromList $ fmap (\s -> (Record s [], 0)) $ tails springs
  44         sigs = M.fromList $ fmap (\g -> (Record [] g, 0)) $ tails signature
  45         lastOperationalChunk =
  46           reverse $ takeWhile isPossOperational $ reverse springs
  47         lastOperational =
  48           M.fromList $ fmap (\s -> (Record s [], 1)) $ tails lastOperationalChunk
  49
  50 fillTableCell, fillTableSigs, fillTable :: Cache -> Record -> Cache
  51 fillTableCell table record
  52   | record `M.member` table = table
  53   | otherwise = M.insert record (opN + signN) table
  54   where (Record springs@(s:ss) signatures@(g:gs)) = record
  55         opN = if (isPossOperational s) then table M.! (Record ss signatures) else 0
  56         signN = if (possibleDamagedChunk springs g) then table M.! (Record (drop (g + 1) springs) gs) else 0
  57
  58 fillTableSigs table (Record springs signatures) = foldr (\gs t -> fillTableCell t (Record springs gs)) table $ tails signatures
  59
  60 fillTable table (Record springs signatures) = foldr (\ss t -> fillTableSigs t (Record ss signatures)) table $ tails springs
  61
  62 isPossOperational :: Spring -> Bool
  63 isPossOperational Operational = True
  64 isPossOperational Unknown = True
  65 isPossOperational _ = False
  66
  67 isPossDamaged :: Spring -> Bool
  68 isPossDamaged Damaged = True
  69 isPossDamaged Unknown = True
  70 isPossDamaged _ = False
  71
  72 possibleDamagedChunk :: [Spring] -> Int -> Bool
  73 possibleDamagedChunk springs n =
  74   isDamagedChunk && ((null afterChunk) || (isPossOperational $ head afterChunk))
  75   where isDamagedChunk = (length $ filter isPossDamaged $ take n springs) == n
  76         afterChunk = drop n springs
  77
  78 -- Parse the input file
  79
  80 recordsP :: Parser [Record]
  81 recordP :: Parser Record
  82 springP :: Parser Spring
  83
  84 recordsP = recordP `sepBy` endOfLine
  85 recordP = Record <$> (many1 springP <* " ") <*> (decimal `sepBy` ",")
  86 springP = (Unknown <$ "?") <|> (Damaged <$ "#") <|> (Operational <$ ".")
  87
  88 successfulParse :: Text -> [Record]
  89 successfulParse input =
  90   case parseOnly recordsP input of
  91     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
  92     Right matches -> matches