Advent of Code 2022 day 16

Day 16 was very similar in style to 2021 day 15 and I reused most of the code from that example. Look to that post for a description of the overall approach, which is A* search again (from day 12) with polymorphism to handle the two types of search.

Representation

I used a fairly direct representation the cave system. A Room is a combination of its valve flow rate and the IDs of its neighbours. A Cave is a Map of room ID to room. A TimedCave is a cave with the search time limit.

type RoomID = String

data Room = Room 
    { _flowRate :: Int
    , _tunnels :: [RoomID]
    } deriving (Eq, Show, Ord)
makeLenses ''Room

type Cave = M.Map RoomID Room
data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int}

I have two types of search state, depending on whether just the person is turning on valves, or both the person and the elephant. Each represents where the person is (or person and and elephant are) and the set of opened valves. Both of these are members of the SearchState typeclass.

data SingleSearchState = SingleSearchState
    { _currentRoom :: RoomID
    , _sOpenValves :: S.Set RoomID
    } deriving (Eq, Show, Ord)
makeLenses ''SingleSearchState

data DoubleSearchState = DoubleSearchState
    { _personRoom :: RoomID
    , _elephantRoom :: RoomID
    , _dOpenValves :: S.Set RoomID
    } deriving (Eq, Show, Ord)
makeLenses ''DoubleSearchState

class (Eq s, Ord s, Show s) => SearchState s where
  emptySearchState :: RoomID -> s
  currentFlow :: s -> CaveContext Int
  successors :: s -> CaveContext (Q.Seq s)
  estimateBenefit :: s -> Int -> CaveContext Int

The search agenda and closed set are both polymorphic on the search state.

data Agendum s = 
    Agendum { _current :: s
            , _trail :: Q.Seq s
            , _trailBenefit :: Int
            , _benefit :: Int
            } deriving (Show, Eq, Ord)
makeLenses ''Agendum   

type Agenda s = P.MaxPQueue Int (Agendum s)

type ExploredStates s = S.Set (s, Int, Int)

Finding next states

The core of handling this domain is generating the next state from a given state. As both the person and the elephant do much the same thing, it made sense to have one function that handled an agent moving. This was a function that was given a current location and set of opened valves, and returned the list of possible (room, valve) pairs the person could occupy next.

personSuccessor, openValveSuccessor, walkSuccessor ::  RoomID -> S.Set RoomID -> CaveContext [(RoomID, S.Set RoomID)]
personSuccessor here opened =
  do ovs <- openValveSuccessor here opened
     ws <- walkSuccessor here opened
     return (ovs ++ ws)

openValveSuccessor here opened
  | here `S.member` opened = return []
  | otherwise = return [(here, S.insert here opened)]

walkSuccessor here opened = 
  do cave <- asks getCave
     let neighbours = (cave ! here) ^. tunnels
     return [(n, opened) | n <- neighbours]

Given those lists, I needed typeclass-specific functions to build the successor search states. For just the person moving, it was simple enough. For both elephant and person moving, I generated all combinations of moves.

-- single person state
successors state = 
  do isFF <- isFullFlow state
     let here = state ^. currentRoom
     let opened = state ^. sOpenValves
     succPairs <- personSuccessor here opened
     let succStates = 
            [ SingleSearchState 
                { _currentRoom = r
                , _sOpenValves = o
                }
            | (r, o) <- succPairs
            ]
     if isFF
     then return $ Q.singleton state
     else return $ Q.fromList succStates

-- double person state
successors state = 
  do isFF <- isFullFlow state
     let pHere = state ^. personRoom
     let eHere = state ^. elephantRoom
     let opened = state ^. dOpenValves
     pSuccPairs <- personSuccessor pHere opened
     eSuccPairs <- personSuccessor eHere opened
     let succStates = 
            [ DoubleSearchState 
                { _personRoom = p
                , _elephantRoom = e
                , _dOpenValves = S.union po eo
                }
            | (p, po) <- pSuccPairs
            , (e, eo) <- eSuccPairs
            ]
     if isFF
     then return $ Q.singleton state
     else return $ Q.fromList succStates

Heuristic

A* search requires a heuristic to guide search. These are often found by relaxing constraints in the problem. In this case, I relaxed the constraint that only certain tunnels existed, leading to a heuristic that turned on the remaining valves, in order, every two minutes.

-- signle person state
estimateBenefit here timeElapsed = 
  do cave <- asks getCave
     timeLimit <- asks getTimeLimit
     let timeRemaining = timeLimit - (timeElapsed + 2)
     cf <- currentFlow here
     let closedValves = (cave `M.withoutKeys` (here ^. sOpenValves)) ^.. folded . flowRate
     let sortedClosedValves = sortOn Down closedValves
     let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
     return $ (cf * timeRemaining) + otherValveFlows

-- double person state
estimateBenefit here timeElapsed = 
  do cave <- asks getCave
     timeLimit <- asks getTimeLimit
     let timeRemaining = timeLimit - (timeElapsed + 2)
     cf <- currentFlow here
     let closedValves = (cave `M.withoutKeys` (here ^. dOpenValves)) ^.. folded . flowRate
     let sortedClosedValves = fmap sum $ chunksOf 2 $ sortOn Down closedValves
     let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
     return $ (cf * timeRemaining) + otherValveFlows

Optimisation

That was just about enough to solve both parts of the puzzle, but I did need two optimisations.

The first was to notice that I may end up turning on all the valves before the time limit. If the cave was a "full flow", with all the valves open, the successor functions returned just one state, which was the same as the original one.

currentFlow state =
  do cave <- asks getCave
     let valves = state ^. dOpenValves
     let presentRooms = cave `M.restrictKeys` valves
     return $ sumOf (folded . flowRate) presentRooms
     
isFullFlow :: SearchState s => s -> CaveContext Bool
isFullFlow state = 
  do cave <- asks getCave
     cf <- currentFlow state
     let ff = sumOf (folded . flowRate) cave
     return (cf == ff)

The other optimisation was to deal with part 2, and the huge search space. The only thing I did there was to restrict the size of the agenda to an arbitrary 5000 elements (good old beam search!). That was enough to find the optimal solution without excessive memory consumption.

Code

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