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.

    Neil Smith

    Read more posts by this author.