advent12/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2023/12/08/advent-of-code-2023-day-8/
   2
   3 import AoC
   4 import Data.Text (Text)
   5 import qualified Data.Text.IO as TIO
   6 import Data.Attoparsec.Text -- hiding (take)
   7 import Control.Applicative
   8 import Data.List
   9 import qualified Data.Map.Strict as M
  10 import Data.Maybe
  11 import Control.Monad.State.Strict
  12
  13 data Spring = Unknown | Damaged | Operational deriving (Show, Eq)
  14 data Record = Record [Spring] [Int] deriving (Show)
  15
  16 type Cache = M.Map Record Int
  17 -- type CacheState = State Cache
  18
  19
  20 main :: IO ()
  21 main =
  22   do  dataFileName <- getDataFileName
  23       text <- TIO.readFile dataFileName
  24       let records = successfulParse text
  25       -- print records
  26       -- print $ fmap numDamagedToPlace records
  27       -- print $ fmap candidates records
  28       -- print $ possibleAssignments (records !! 1)
  29       -- print $ fmap countViableAssignments records
  30       print $ part1 records
  31       print $ unfoldRecord $ head records
  32       print $ part2 records
  33
  34 part1, part2 :: [Record] -> Int
  35 part1 = sum . fmap countViableAssignments
  36 part2 = sum . fmap (countViableAssignments . unfoldRecord)
  37
  38 unfoldRecord :: Record -> Record
  39 unfoldRecord (Record springs signature) = Record uSprings uSignature
  40   where uSprings = intercalate [Unknown] $ replicate 5 springs
  41         uSignature = concat $ replicate 5 signature
  42
  43
  44 countViable :: Record -> CacheState Int
  45
  46
  47 countViable previousSprings (Record (s:springs) (g:signature)) =
  48
  49
  50 -- if next few springs can match next group:
  51 --   countViable (springs after group) (tail groups) + countViable (tail springs) groups
  52 -- else
  53 --   countViable (tail springs) groups
  54
  55
  56
  57 countViableAssignments :: Record -> Int
  58 countViableAssignments = length . filter matchesSignature . possibleAssignments
  59
  60 matchesSignature :: Record -> Bool
  61 matchesSignature (Record springs signature) = signSprings springs == signature
  62
  63 signSprings :: [Spring] -> [Int]
  64 signSprings = fmap (length) . filter ((== Damaged) . head) . group
  65
  66 choose :: Int -> [a] -> [[a]]
  67 choose 0 _ = [[]]
  68 choose n (x:xs)
  69   | length xs == n - 1 = [(x:xs)]
  70   | otherwise = (fmap (x:) (choose (n-1) xs)) ++ (choose n xs)
  71
  72 -- unknownIndices :: [Spring] -> [Int]
  73 -- unknownIndices = elemIndices Unknown
  74
  75 numDamagedToPlace :: Record -> Int
  76 numDamagedToPlace (Record springs signature) = totalDamaged - knownDamaged
  77   where knownDamaged = length $ filter (== Damaged) springs
  78         totalDamaged = sum signature
  79
  80 candidates :: Record -> [[Int]]
  81 candidates r@(Record springs _) =
  82   choose (numDamagedToPlace r) (elemIndices Unknown springs)
  83
  84 replaceUnknowns :: [Spring] -> [Int] -> [Spring]
  85 replaceUnknowns springs indices = foldr go [] indexedSprings
  86   where indexedSprings = zip [0..] springs
  87         go (i, Unknown) acc = if (i `elem` indices) then Damaged:acc
  88                                                     else Operational:acc
  89         go (_, s) acc = s:acc
  90
  91 possibleAssignments :: Record -> [Record]
  92 possibleAssignments r@(Record springs signature) =
  93   fmap (\p -> Record p signature) possibles
  94   where cands = candidates r
  95         possibles = fmap (replaceUnknowns springs) cands
  96
  97 -- Parse the input file
  98
  99 recordsP :: Parser [Record]
 100 recordP :: Parser Record
 101 springP :: Parser Spring
 102
 103 recordsP = recordP `sepBy` endOfLine
 104 recordP = Record <$> (many1 springP <* " ") <*> (decimal `sepBy` ",")
 105 springP = (Unknown <$ "?") <|> (Damaged <$ "#") <|> (Operational <$ ".")
 106
 107 successfulParse :: Text -> [Record]
 108 successfulParse input =
 109   case parseOnly recordsP input of
 110     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 111     Right matches -> matches