X-Git-Url: https://git.njae.me.uk/?a=blobdiff_plain;f=advent12%2FMain.hs;h=c0228990a9c89f1018f7fd21f0bf9e51a402f4e8;hb=c42805c934d8b0300b8c4d9ba10b1162585b31b7;hp=59d73a365022b40c31bff09ddaa88420b4d498c6;hpb=8761302ae0a2266a934a49138e4c01f3fa5e4af1;p=advent-of-code-21.git

diff --git a/advent12/Main.hs b/advent12/Main.hs
index 59d73a3..c022899 100644
--- a/advent12/Main.hs
+++ b/advent12/Main.hs
@@ -1,5 +1,4 @@
--- Writeup at https://work.njae.me.uk/2021/12/09/advent-of-code-2021-day-8/
-
+-- Writeup at https://work.njae.me.uk/2021/12/13/advent-of-code-2021-day-12/
 
 import Data.Text ()
 import qualified Data.Text.IO as TIO
@@ -8,38 +7,41 @@ import Data.Attoparsec.Text
 -- import Control.Applicative
 
 import Data.Tuple
--- import Data.List
 import Data.Char
 import qualified Data.Map.Strict as M
 import Data.Map.Strict ((!))
 import qualified Data.Set as S
 import Data.Set ((\\))
 
-data Path = Path String [String] (S.Set String)
+data Path = Path String         -- current cave
+                 [String]       -- caves visited
+                 (S.Set String) -- closed set of small caves visited
+                 (Maybe String) -- the small cave we've visited twice
   deriving (Eq, Ord, Show)
 
 type PathSet = S.Set Path
 type Graph = M.Map String (S.Set String)
 
 
-
 main :: IO ()
 main = 
   do  text <- TIO.readFile "data/advent12.txt"
       let edges = successfulParse text
-      print edges
       let graph = mkGraph edges
-      print graph
-      print $ part1 graph
-      -- print $ part1 displays
-      -- print $ part2 displays
+      let paths = allPaths graph (S.singleton (Path "start" [] S.empty Nothing)) S.empty
+      print $ part1 paths
+      print $ part2 paths
 
 mkGraph :: [(String, String)] -> Graph
 mkGraph edges = foldr mkEdge pass1 $ map swap edges
   where pass1 = foldr mkEdge M.empty edges
-        mkEdge (here, there) g = M.insertWith (S.union) here (S.singleton there) g
+        mkEdge (here, there) = M.insertWith (S.union) here (S.singleton there)
+
+part1 :: PathSet -> Int
+part1 paths = S.size $ S.filter nonReturning paths
 
-part1 graph = S.size $ allPaths graph (S.singleton (Path "start" [] S.empty)) S.empty
+part2 :: PathSet -> Int
+part2 paths = S.size paths
 
 allPaths :: Graph -> PathSet -> PathSet -> PathSet
 allPaths graph agenda results
@@ -51,29 +53,37 @@ allPaths graph agenda results
         results' = S.union results $ recordResult current
 
 extendPath :: Graph -> Path -> PathSet
-extendPath graph (Path current trail visited) 
+extendPath graph (Path current trail visited returned) 
   | current == "end" = S.empty
-  | otherwise = S.map newPath visitable
+  | (current == "start") && (current `S.member` visited) = S.empty
+  | otherwise = S.union (S.map newPathNovel visitableNovel) 
+                        (S.map newPathReturning visitableReturning)
   where neighbours = graph ! current
         visited' = if isSmall current then S.insert current visited else visited
         trail' = (current:trail)
-        visitable = neighbours \\ visited
-        newPath next = Path next trail' visited'
+        visitableNovel = neighbours \\ visited -- if we're not returning to a small cave
+        visitableReturning = if returned == Nothing 
+          then (S.filter isSmall neighbours) `S.intersection` visited -- returning to a small cave already visited
+          else S.empty
+        newPathNovel next = Path next trail' visited' returned
+        newPathReturning next = Path next trail' visited' (Just next)
 
 recordResult :: Path -> PathSet
-recordResult path@(Path current _trail _visited)
-  | current == "end" = S.singleton path -- (Path current trail visited)
+recordResult path@(Path current _ _ _)
+  | current == "end" = S.singleton path
   | otherwise = S.empty
 
-
 isSmall :: String -> Bool
 isSmall = all isLower
-isBig = not . isSmall
+
+nonReturning :: Path -> Bool
+nonReturning (Path _ _ _ Nothing) = True
+nonReturning (Path _ _ _ (Just _)) = False
 
 -- Parse the input file
 
 graphP = edgeP `sepBy` endOfLine
-edgeP = (,) <$> (many1 letter <* "-") <*> many1 letter
+edgeP = (,) <$> many1 letter <* "-" <*> many1 letter
 
 -- successfulParse :: Text -> (Integer, [Maybe Integer])
 successfulParse input =