advent12/Main.hs

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