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.List
  12 import Data.Char
  13 import qualified Data.Map.Strict as M
  14 import Data.Map.Strict ((!))
  15 import qualified Data.Set as S
  16 import Data.Set ((\\))
  17
  18 data Path = Path String [String] (S.Set String)
  19   deriving (Eq, Ord, Show)
  20
  21 type PathSet = S.Set Path
  22 type Graph = M.Map String (S.Set String)
  23
  24
  25
  26 main :: IO ()
  27 main =
  28   do  text <- TIO.readFile "data/advent12.txt"
  29       let edges = successfulParse text
  30       print edges
  31       let graph = mkGraph edges
  32       print graph
  33       print $ part1 graph
  34       -- print $ part1 displays
  35       -- print $ part2 displays
  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) g = M.insertWith (S.union) here (S.singleton there) g
  41
  42 part1 graph = S.size $ allPaths graph (S.singleton (Path "start" [] S.empty)) S.empty
  43
  44 allPaths :: Graph -> PathSet -> PathSet -> PathSet
  45 allPaths graph agenda results
  46   | S.null agenda = results
  47   | otherwise = allPaths graph agenda'' results'
  48   where (current, agenda') = S.deleteFindMin agenda
  49         newPaths = extendPath graph current
  50         agenda'' = S.union agenda' newPaths
  51         results' = S.union results $ recordResult current
  52
  53 extendPath :: Graph -> Path -> PathSet
  54 extendPath graph (Path current trail visited)
  55   | current == "end" = S.empty
  56   | otherwise = S.map newPath visitable
  57   where neighbours = graph ! current
  58         visited' = if isSmall current then S.insert current visited else visited
  59         trail' = (current:trail)
  60         visitable = neighbours \\ visited
  61         newPath next = Path next trail' visited'
  62
  63 recordResult :: Path -> PathSet
  64 recordResult path@(Path current _trail _visited)
  65   | current == "end" = S.singleton path -- (Path current trail visited)
  66   | otherwise = S.empty
  67
  68
  69 isSmall :: String -> Bool
  70 isSmall = all isLower
  71 isBig = not . isSmall
  72
  73 -- Parse the input file
  74
  75 graphP = edgeP `sepBy` endOfLine
  76 edgeP = (,) <$> (many1 letter <* "-") <*> many1 letter
  77
  78 -- successfulParse :: Text -> (Integer, [Maybe Integer])
  79 successfulParse input =
  80   case parseOnly graphP input of
  81     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
  82     Right graph -> graph