advent12/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2021/12/13/advent-of-code-2021-day-12/
   2
   3 import Data.Text ()
   4 import qualified Data.Text.IO as TIO
   5
   6 import Data.Attoparsec.Text
   7 -- import Control.Applicative
   8
   9 import Data.Tuple
  10 import Data.Char
  11 import qualified Data.Map.Strict as M
  12 import Data.Map.Strict ((!))
  13 import qualified Data.Set as S
  14 import Data.Set ((\\))
  15
  16 data Path = Path String         -- current cave
  17                  [String]       -- caves visited
  18                  (S.Set String) -- closed set of small caves visited
  19                  (Maybe String) -- the small cave we've visited twice
  20   deriving (Eq, Ord, Show)
  21
  22 type PathSet = S.Set Path
  23 type Graph = M.Map String (S.Set String)
  24
  25
  26 main :: IO ()
  27 main =
  28   do  text <- TIO.readFile "data/advent12.txt"
  29       let edges = successfulParse text
  30       let graph = mkGraph edges
  31       let paths = allPaths graph (S.singleton (Path "start" [] S.empty Nothing)) S.empty
  32       print $ part1 paths
  33       print $ part2 paths
  34
  35 mkGraph :: [(String, String)] -> Graph
  36 mkGraph edges = foldr mkEdge pass1 $ map swap edges
  37   where pass1 = foldr mkEdge M.empty edges
  38         mkEdge (here, there) = M.insertWith (S.union) here (S.singleton there)
  39
  40 part1 :: PathSet -> Int
  41 part1 paths = S.size $ S.filter nonReturning paths
  42
  43 part2 :: PathSet -> Int
  44 part2 paths = S.size paths
  45
  46 allPaths :: Graph -> PathSet -> PathSet -> PathSet
  47 allPaths graph agenda results
  48   | S.null agenda = results
  49   | otherwise = allPaths graph agenda'' results'
  50   where (current, agenda') = S.deleteFindMin agenda
  51         newPaths = extendPath graph current
  52         agenda'' = S.union agenda' newPaths
  53         results' = S.union results $ recordResult current
  54
  55 extendPath :: Graph -> Path -> PathSet
  56 extendPath graph (Path current trail visited returned)
  57   | current == "end" = S.empty
  58   | (current == "start") && (current `S.member` visited) = S.empty
  59   | otherwise = S.union (S.map newPathNovel visitableNovel)
  60                         (S.map newPathReturning visitableReturning)
  61   where neighbours = graph ! current
  62         visited' = if isSmall current then S.insert current visited else visited
  63         trail' = (current:trail)
  64         visitableNovel = neighbours \\ visited -- if we're not returning to a small cave
  65         visitableReturning = if returned == Nothing
  66           then (S.filter isSmall neighbours) `S.intersection` visited -- returning to a small cave already visited
  67           else S.empty
  68         newPathNovel next = Path next trail' visited' returned
  69         newPathReturning next = Path next trail' visited' (Just next)
  70
  71 recordResult :: Path -> PathSet
  72 recordResult path@(Path current _ _ _)
  73   | current == "end" = S.singleton path
  74   | otherwise = S.empty
  75
  76 isSmall :: String -> Bool
  77 isSmall = all isLower
  78
  79 nonReturning :: Path -> Bool
  80 nonReturning (Path _ _ _ Nothing) = True
  81 nonReturning (Path _ _ _ (Just _)) = False
  82
  83 -- Parse the input file
  84
  85 graphP = edgeP `sepBy` endOfLine
  86 edgeP = (,) <$> many1 letter <* "-" <*> many1 letter
  87
  88 -- successfulParse :: Text -> (Integer, [Maybe Integer])
  89 successfulParse input =
  90   case parseOnly graphP input of
  91     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
  92     Right graph -> graph