Done day 12 part 2

diff --git a/advent12/Main.hs b/advent12/Main.hs
index 59d73a365022b40c31bff09ddaa88420b4d498c6..a191524eaeee1810f59c0562545b8ecb9c7ef09d 100644 (file)
--- a/advent12/Main.hs
+++ b/advent12/Main.hs
@@ -8,14 +8,16 @@ 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 cavees visited
+                 (Maybe String) -- the small cave we've visited twice
   deriving (Eq, Ord, Show)
 
 type PathSet = S.Set Path
@@ -27,19 +29,21 @@ 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 graph = S.size $ allPaths graph (S.singleton (Path "start" [] S.empty)) S.empty
+part1 :: PathSet -> Int
+part1 paths = S.size $ S.filter nonReturning paths
+
+part2 :: PathSet -> Int
+part2 paths = S.size paths
 
 allPaths :: Graph -> PathSet -> PathSet -> PathSet
 allPaths graph agenda results
@@ -51,29 +55,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)
+recordResult path@(Path current _ _ _)
   | current == "end" = S.singleton path -- (Path current trail visited)
   | 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 =