advent12/MainBruteForce.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
  10 data Spring = Unknown | Damaged | Operational deriving (Show, Eq)
  11 data Record = Record [Spring] [Int] deriving (Show)
  12
  13 main :: IO ()
  14 main =
  15   do  dataFileName <- getDataFileName
  16       text <- TIO.readFile dataFileName
  17       let records = successfulParse text
  18       -- print records
  19       -- print $ fmap numDamagedToPlace records
  20       -- print $ fmap candidates records
  21       -- print $ possibleAssignments (records !! 1)
  22       -- print $ fmap countViableAssignments records
  23       print $ part1 records
  24       -- print $ unfoldRecord $ head records
  25       -- print $ part2 records
  26
  27 part1, part2 :: [Record] -> Int
  28 part1 = sum . fmap countViableAssignments
  29 part2 = sum . fmap (countViableAssignments . unfoldRecord)
  30
  31 unfoldRecord :: Record -> Record
  32 unfoldRecord (Record springs signature) = Record uSprings uSignature
  33   where uSprings = intercalate [Unknown] $ replicate 5 springs
  34         uSignature = concat $ replicate 5 signature
  35
  36
  37 countViableAssignments :: Record -> Int
  38 countViableAssignments = length . filter matchesSignature . possibleAssignments
  39
  40 matchesSignature :: Record -> Bool
  41 matchesSignature (Record springs signature) = signSprings springs == signature
  42
  43 signSprings :: [Spring] -> [Int]
  44 signSprings = fmap (length) . filter ((== Damaged) . head) . group
  45
  46 choose :: Int -> [a] -> [[a]]
  47 choose 0 _ = [[]]
  48 choose n (x:xs)
  49   | length xs == n - 1 = [(x:xs)]
  50   | otherwise = (fmap (x:) (choose (n-1) xs)) ++ (choose n xs)
  51
  52 -- unknownIndices :: [Spring] -> [Int]
  53 -- unknownIndices = elemIndices Unknown
  54
  55 numDamagedToPlace :: Record -> Int
  56 numDamagedToPlace (Record springs signature) = totalDamaged - knownDamaged
  57   where knownDamaged = length $ filter (== Damaged) springs
  58         totalDamaged = sum signature
  59
  60 candidates :: Record -> [[Int]]
  61 candidates r@(Record springs _) =
  62   choose (numDamagedToPlace r) (elemIndices Unknown springs)
  63
  64 replaceUnknowns :: [Spring] -> [Int] -> [Spring]
  65 replaceUnknowns springs indices = foldr go [] indexedSprings
  66   where indexedSprings = zip [0..] springs
  67         go (i, Unknown) acc = if (i `elem` indices) then Damaged:acc
  68                                                     else Operational:acc
  69         go (_, s) acc = s:acc
  70
  71 possibleAssignments :: Record -> [Record]
  72 possibleAssignments r@(Record springs signature) =
  73   fmap (\p -> Record p signature) possibles
  74   where cands = candidates r
  75         possibles = fmap (replaceUnknowns springs) cands
  76
  77 -- Parse the input file
  78
  79 recordsP :: Parser [Record]
  80 recordP :: Parser Record
  81 springP :: Parser Spring
  82
  83 recordsP = recordP `sepBy` endOfLine
  84 recordP = Record <$> (many1 springP <* " ") <*> (decimal `sepBy` ",")
  85 springP = (Unknown <$ "?") <|> (Damaged <$ "#") <|> (Operational <$ ".")
  86
  87 successfulParse :: Text -> [Record]
  88 successfulParse input =
  89   case parseOnly recordsP input of
  90     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
  91     Right matches -> matches