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, takeWhile)
   7 import Control.Applicative
   8 import Data.List
   9 import qualified Data.Map.Strict as M
  10 import Data.Maybe
  11
  12 data Spring = Unknown | Damaged | Operational deriving (Show, Eq, Ord)
  13 data Record = Record [Spring] [Int] deriving (Show, Eq, Ord)
  14
  15 type Cache = M.Map Record Int
  16 -- type CacheState = State Cache
  17
  18
  19 main :: IO ()
  20 main =
  21   do  dataFileName <- getDataFileName
  22       text <- TIO.readFile dataFileName
  23       let records = successfulParse text
  24       -- print records
  25       -- print $ fmap numDamagedToPlace records
  26       -- print $ fmap candidates records
  27       -- print $ possibleAssignments (records !! 1)
  28       -- print $ fmap countViableAssignments records
  29       print $ part1 records
  30       print $ unfoldRecord $ head records
  31       -- print $ part2 records
  32
  33 part1 :: [Record] -> Int
  34 part1 = sum . fmap countViableAssignments
  35 -- part2 = sum . fmap (countViableAssignments . unfoldRecord)
  36
  37 unfoldRecord :: Record -> Record
  38 unfoldRecord (Record springs signature) = Record uSprings uSignature
  39   where uSprings = intercalate [Unknown] $ replicate 5 springs
  40         uSignature = concat $ replicate 5 signature
  41
  42
  43
  44 initialCache :: Record -> Cache
  45 initialCache (Record springs signature) = M.union lastOperational cache0
  46   where cache0 = M.union sprs sigs
  47         sprs = M.fromList $ fmap (\s -> (Record s [], 0)) $ tails springs
  48         sigs = M.fromList $ fmap (\g -> (Record [] g, 0)) $ tails signature
  49         lastOperationalChunk =
  50           reverse $ takeWhile isPossOperational $ reverse springs
  51         lastOperational =
  52           M.fromList $ fmap (\s -> (Record s [], 1)) $ tails lastOperationalChunk
  53
  54 isPossOperational :: Spring -> Bool
  55 isPossOperational Operational = True
  56 isPossOperational Unknown = True
  57 isPossOperational _ = False
  58
  59 isPossDamaged :: Spring -> Bool
  60 isPossDamaged Damaged = True
  61 isPossDamaged Unknown = True
  62 isPossDamaged _ = False
  63
  64 possibleDamagedChunk :: [Spring] -> Int -> Bool
  65 possibleDamagedChunk springs n =
  66   isDamagedChunk && ((null afterChunk) || (possOperational $ head afterChunk))
  67   where isDamagedChunk = (length $ filter isPossDamaged $ take n springs) == n
  68         afterChunk = take 1 $ drop n springs
  69
  70
  71 -- countViable previousSprings (Record (s:springs) (g:signature)) =
  72
  73
  74 -- count: either consume this group or not
  75
  76
  77 -- cache is how many ways to assign unknowns, leaving this partial record.
  78 -- first item of springs is unknown or damaged
  79 -- count = count_consuming_group + count_leaving_group
  80 -- if first spring is damaged: count = count_consuming_group
  81
  82
  83
  84 -- if first spring is damaged and next few springs can match next group:
  85 --   consume springs (including all following operational ones), consume group
  86 --   Add to cache: record' -> cache ! original
  87 --   return countViable record'
  88 -- if first spring is unknown
  89 --   assume it's damaged, consume springs, consume group
  90
  91
  92
  93
  94
  95 --   countViable (springs after group) (tail groups) + countViable (tail springs) groups
  96 -- else
  97 --   countViable (tail springs) groups
  98
  99
 100
 101 countViableAssignments :: Record -> Int
 102 countViableAssignments = length . filter matchesSignature . possibleAssignments
 103
 104 matchesSignature :: Record -> Bool
 105 matchesSignature (Record springs signature) = signSprings springs == signature
 106
 107 signSprings :: [Spring] -> [Int]
 108 signSprings = fmap (length) . filter ((== Damaged) . head) . group
 109
 110 choose :: Int -> [a] -> [[a]]
 111 choose 0 _ = [[]]
 112 choose n (x:xs)
 113   | length xs == n - 1 = [(x:xs)]
 114   | otherwise = (fmap (x:) (choose (n-1) xs)) ++ (choose n xs)
 115
 116 -- unknownIndices :: [Spring] -> [Int]
 117 -- unknownIndices = elemIndices Unknown
 118
 119 numDamagedToPlace :: Record -> Int
 120 numDamagedToPlace (Record springs signature) = totalDamaged - knownDamaged
 121   where knownDamaged = length $ filter (== Damaged) springs
 122         totalDamaged = sum signature
 123
 124 candidates :: Record -> [[Int]]
 125 candidates r@(Record springs _) =
 126   choose (numDamagedToPlace r) (elemIndices Unknown springs)
 127
 128 replaceUnknowns :: [Spring] -> [Int] -> [Spring]
 129 replaceUnknowns springs indices = foldr go [] indexedSprings
 130   where indexedSprings = zip [0..] springs
 131         go (i, Unknown) acc = if (i `elem` indices) then Damaged:acc
 132                                                     else Operational:acc
 133         go (_, s) acc = s:acc
 134
 135 possibleAssignments :: Record -> [Record]
 136 possibleAssignments r@(Record springs signature) =
 137   fmap (\p -> Record p signature) possibles
 138   where cands = candidates r
 139         possibles = fmap (replaceUnknowns springs) cands
 140
 141 -- Parse the input file
 142
 143 recordsP :: Parser [Record]
 144 recordP :: Parser Record
 145 springP :: Parser Spring
 146
 147 recordsP = recordP `sepBy` endOfLine
 148 recordP = Record <$> (many1 springP <* " ") <*> (decimal `sepBy` ",")
 149 springP = (Unknown <$ "?") <|> (Damaged <$ "#") <|> (Operational <$ ".")
 150
 151 successfulParse :: Text -> [Record]
 152 successfulParse input =
 153   case parseOnly recordsP input of
 154     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 155     Right matches -> matches