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.