cleaner types for day 12 before refactor algorithm

Author: mstksg

src/AOC/Challenge/Day12.hs

@@ -28,6 +28,9 @@ module AOC.Challenge.Day12 (
 
 import AOC.Prelude
 
+import Data.Bitraversable
+import Data.Finitary
+import qualified AOC.Common.FinitarySet as FS
 import qualified Data.Graph.Inductive as G
 import qualified Data.IntMap as IM
 import qualified Data.IntSet as IS
@@ -45,48 +48,114 @@ import qualified Text.Megaparsec as P
 import qualified Text.Megaparsec.Char as P
 import qualified Text.Megaparsec.Char.Lexer as PP
 
--- False: small
--- True: large
-toAdjMatrix :: [(String, String)] -> Map String [(Bool, String)]
-toAdjMatrix = fmap (mapMaybe postProcess)
- . M.fromListWith (++)
+data SrcNode = StartNode
+ | SrcCave Cave
+ deriving stock (Generic, Eq, Ord, Show)
+ deriving anyclass (Finitary)
+
+data DestNode = DestCave Cave
+ | EndNode
+ deriving stock (Generic, Eq, Ord, Show)
+ deriving anyclass (Finitary)
+
+data FullNode = FullStart
+ | FullCave Cave
+ | FullEnd
+ deriving stock (Generic, Eq, Ord, Show)
+ deriving anyclass (Finitary)
+
+data Ident = Ident (Finite 26) (Maybe (Finite 26))
+ deriving stock (Generic, Eq, Ord, Show)
+ deriving anyclass (Finitary)
+
+data Cave = Big Ident
+ | Small Ident
+ deriving stock (Generic, Eq, Ord, Show)
+ deriving anyclass (Finitary)
+
+classify :: String -> Maybe FullNode
+classify "start" = Just FullStart
+classify "end" = Just FullEnd
+classify [a] = charFinite a <&> \case
+ (False, x) -> FullCave $ Small (Ident x Nothing)
+ (True , x) -> FullCave $ Big (Ident x Nothing)
+classify [a,b] = do
+ (cx, x) <- charFinite a
+ (cy, y) <- charFinite b
+ let ident = Ident x (Just y)
+ case (cx, cy) of
+ (False, False) -> Just . FullCave $ Small ident
+ (True , True ) -> Just . FullCave $ Big ident
+ _ -> Nothing
+classify _ = Nothing
+
+srcDest :: FullNode -> (Maybe SrcNode, Maybe DestNode)
+srcDest = \case
+ FullStart -> (Just StartNode, Nothing)
+ FullCave c -> (Just (SrcCave c), Just (DestCave c))
+ FullEnd -> (Nothing, Just EndNode)
+
+toAdjMatrix :: [(FullNode, FullNode)] -> Map SrcNode [DestNode]
+toAdjMatrix = M.fromListWith (++)
  . concatMap (uncurry buildLinks)
  where
- buildLinks a b = [(a, [b]), (b, [a])]
- postProcess str = (all isUpper str, str) <$ guard (str /= "start")
+ toLink a b = (a, [b])
+ buildLinks a b = catMaybes
+ [ toLink <$> srcA <*> destB
+ , toLink <$> srcB <*> destA
+ ]
+ where
+ (srcA, destA) = srcDest a
+ (srcB, destB) = srcDest b
 
-day12a :: [(String, String)] :~> Int
+day12a :: [(FullNode, FullNode)] :~> Int
 day12a = MkSol
- { sParse = traverseLines $ listTup . splitOn "-"
+ { sParse = traverseLines $ bitraverse classify classify <=< listTup . splitOn "-"
  , sShow = show
  , sSolve = Just . length . findPaths . toAdjMatrix
  }
 
-findPaths :: Map String [(Bool, String)] -> [[String]]
-findPaths mp = go S.empty "start"
+findPaths :: Map SrcNode [DestNode] -> [[Cave]]
+findPaths mp = do
+ nextBranch <- mp M.! StartNode
+ case nextBranch of
+ EndNode -> pure []
+ DestCave v@(Big _) -> go FS.empty v
+ DestCave v@(Small c) -> go (FS.singleton c) v
  where
- go seen currPos
- | currPos == "end" = pure ["end"]
- | otherwise = do
- (isLarge, nextBranch) <- mp M.! currPos
- guard $ isLarge || (nextBranch `S.notMember` seen)
- (currPos:) <$> go (S.insert nextBranch seen) nextBranch
-
-day12b :: [(String, String)] :~> Int
+ go :: FS.FinitarySet Ident -> Cave -> [[Cave]]
+ go seen currPos = do
+ nextBranch <- mp M.! SrcCave currPos
+ case nextBranch of
+ EndNode -> pure [currPos]
+ DestCave v@(Big _) -> (currPos:) <$> go seen v
+ DestCave v@(Small c) -> do
+ guard $ c `FS.notMember` seen
+ (currPos:) <$> go (c `FS.insert` seen) v
+
+day12b :: [(FullNode, FullNode)] :~> Int
 day12b = MkSol
  { sParse = sParse day12a
  , sShow = show
  , sSolve = Just . length . findPaths2 . toAdjMatrix
  }
 
-findPaths2 :: Map String [(Bool, String)] -> [[String]]
-findPaths2 mp = go S.empty Nothing "start"
+findPaths2 :: Map SrcNode [DestNode] -> [[Cave]]
+findPaths2 mp = do
+ nextBranch <- mp M.! StartNode
+ case nextBranch of
+ EndNode -> pure []
+ DestCave v@(Big _) -> go FS.empty Nothing v
+ DestCave v@(Small c) -> go (FS.singleton c) Nothing v
  where
- go seen seenTwice currPos
- | currPos == "end" = pure ["end"]
- | otherwise = do
- (isLarge, nextBranch) <- mp M.! currPos
- newSeenTwice <- if not isLarge && (nextBranch `S.member` seen)
- then Just nextBranch <$ guard (isNothing seenTwice)
+ go :: FS.FinitarySet Ident -> Maybe Ident -> Cave -> [[Cave]]
+ go seen seenTwice currPos = do
+ nextBranch <- mp M.! SrcCave currPos
+ case nextBranch of
+ EndNode -> pure [currPos]
+ DestCave v@(Big _) -> (currPos:) <$> go seen seenTwice v
+ DestCave v@(Small c) -> do
+ newSeenTwice <- if c `FS.member` seen
+ then Just c <$ guard (isNothing seenTwice)
  else pure seenTwice
- (currPos:) <$> go (S.insert nextBranch seen) newSeenTwice nextBranch
+ (currPos:) <$> go (c `FS.insert` seen) newSeenTwice v