Advent of Code 2021 day 12

Day 12 had some fiddly bits, so things ended up fairly complex. In essence, the problem is an exhaustive search for paths on a graph, but the rules on revisiting nodes makes for complex code. In addition, the extension task for part 2 was such that it ended up easier to solve part 2 first, then extract just the solutions that were valid for part 1.

As with all search-type tasks, I find it easier to think about the problem as having an agenda of partial paths, with the agenda being updated as the search progresses.

In this case, an agendum was represented as a Path containing the current cave, the history of all caves visited, the closed set of small caves visited, and an optional record of a small cave that was visited twice.

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)

The allPaths function finds all the paths in the cave complex. It takes as input the graph (the cave complex), a Set of partial Paths, and a Set of the completed routes found so far. It extracts an arbitrary element of the agenda, extends it if possible, then adds the extended paths back into the agenda. If the path ends at the end cave, the path is added to the set of results.

allPaths :: Graph -> PathSet -> PathSet -> PathSet
allPaths graph agenda results
  | S.null agenda = results
  | otherwise = allPaths graph agenda'' results'
  where (current, agenda') = S.deleteFindMin agenda
        newPaths = extendPath graph current
        agenda'' = S.union agenda' newPaths
        results' = S.union results $ recordResult current

recordResult :: Path -> PathSet
recordResult path@(Path current _ _ _)
  | current == "end" = S.singleton path
  | otherwise = S.empty

All the tricky details are handled in extendPath. It's given a Path and returns a Set of new Paths that lead on from it. There are a lot of things to check, so the logic here needs some explanation.

1. If the path gets as far as "end", there are no extensions.
2. If the path has returned to "start" (the current cave is "start" and "start" is in the visited set), there are no extensions.
3. Othewise, there are extensions.

There are two types of extension:

1. Visting a cave not in the visited closed set
2. Visiting a cave in the visited set, but we've not revisited a small cave yet.

To handle the first case, I find all the neighbours of the current cave, remove all the ones in visited, and create a new Path for each.

To handle the second case, I check that we've not already used our "return to a small cave" exemption (returned == Nothing), then extend the path for every reachable small cave we've already visited.

In both cases, I update visited if I'm currently in a small cave.

It's a fiddly mess.

extendPath :: Graph -> Path -> PathSet
extendPath graph (Path current trail visited returned) 
  | current == "end" = S.empty
  | (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)
        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)

That's the solution to part 2. The solution to part 1 is the same, only without the paths that take advantage of the "revisit a small cave" allowance.

main :: IO ()
main = 
  do  text <- TIO.readFile "data/advent12.txt"
      let edges = successfulParse text
      let graph = mkGraph edges
      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) = M.insertWith (S.union) here (S.singleton there)

part1 :: PathSet -> Int
part1 paths = S.size $ S.filter nonReturning paths

part2 :: PathSet -> Int
part2 paths = S.size paths


nonReturning :: Path -> Bool
nonReturning (Path _ _ _ Nothing) = True
nonReturning (Path _ _ _ (Just _)) = False

Code

You can get the code from my locally-hosed Git repo, or from Gitlab.