--- /dev/null
+-- Writeup at https://work.njae.me.uk/2022/12/17/advent-of-code-2022-day-16/
+
+-- import Debug.Trace
+
+import AoC
+import Data.Text (Text)
+import qualified Data.Text.IO as TIO
+import Data.Attoparsec.Text hiding (take, D)
+import Control.Applicative
+import qualified Data.PQueue.Prio.Max as P
+import qualified Data.Set as S
+import qualified Data.Sequence as Q
+import qualified Data.Map.Strict as M
+import Data.Map.Strict ((!))
+import Data.Sequence ((|>))
+import Data.List
+import Data.List.Split (chunksOf)
+import Data.Ord
+import Control.Monad.Reader
+import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
+
+-- pattern Empty <- (Q.viewl -> Q.EmptyL) where Empty = Q.empty
+-- pattern x :< xs <- (Q.viewl -> x Q.:< xs) where (:<) = (Q.<|)
+-- pattern xs :> x <- (Q.viewr -> xs Q.:> x) where (:>) = (Q.|>)
+
+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}
+
+type CaveContext = Reader TimedCave
+
+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
+
+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)
+
+
+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
+
+instance SearchState SingleSearchState where
+ emptySearchState startID = SingleSearchState { _currentRoom = startID, _sOpenValves = S.empty }
+
+ currentFlow state =
+ do cave <- asks getCave
+ let valves = state ^. sOpenValves
+ let presentRooms = cave `M.restrictKeys` valves
+ return $ sumOf (folded . flowRate) presentRooms
+
+ 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
+
+ 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
+
+
+instance SearchState DoubleSearchState where
+ emptySearchState startID = DoubleSearchState
+ { _personRoom = startID
+ , _elephantRoom = startID
+ , _dOpenValves = S.empty
+ }
+
+ currentFlow state =
+ do cave <- asks getCave
+ let valves = state ^. dOpenValves
+ let presentRooms = cave `M.restrictKeys` valves
+ return $ sumOf (folded . flowRate) presentRooms
+
+ 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
+
+ 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
+
+
+main :: IO ()
+main =
+ do dataFileName <- getDataFileName
+ text <- TIO.readFile dataFileName
+ let cave = successfulParse text
+ -- print cave
+ print $ part1 cave
+ print $ part2 cave
+
+-- part1 :: Cave -> Maybe (Agendum SingleSearchState)
+-- part1 cave = runReader (searchCave "AA") (TimedCave cave 30)
+
+-- part2 :: Cave -> Maybe (Agendum DoubleSearchState)
+-- part2 cave = runReader (searchCave "AA") (TimedCave cave 26)
+
+part1, part2 :: Cave -> Int
+part1 cave = maybe 0 _benefit result
+ where result = runReader (searchCave "AA") (TimedCave cave 30) :: Maybe (Agendum SingleSearchState)
+part2 cave = maybe 0 _benefit result
+ where result = runReader (searchCave "AA") (TimedCave cave 26) :: Maybe (Agendum DoubleSearchState)
+
+searchCave :: SearchState s => String -> CaveContext (Maybe (Agendum s))
+searchCave startRoom =
+ do agenda <- initAgenda startRoom
+ aStar agenda S.empty
+
+initAgenda :: SearchState s => String -> CaveContext (Agenda s)
+initAgenda startID =
+ do let startState = emptySearchState startID
+ b <- estimateBenefit startState 0
+ return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
+
+aStar :: SearchState s => Agenda s -> ExploredStates s -> CaveContext (Maybe (Agendum s))
+aStar agenda closed
+ -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
+ -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : len " ++ (show $ Q.length $ _trail $ snd $ P.findMax agenda)) False = undefined
+ | P.null agenda = return Nothing
+ | otherwise =
+ do let (_, currentAgendum) = P.findMax agenda
+ let reached = currentAgendum ^. current
+ nexts <- candidates currentAgendum closed
+ let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
+ -- let beamAgenda = P.fromDescList $ P.take 10000 newAgenda -- agenda beam width
+ let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
+ reachedGoal <- isGoal currentAgendum
+ let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
+ if reachedGoal
+ then return (Just currentAgendum)
+ else if (cl `S.member` closed)
+ then aStar (P.deleteMax agenda) closed
+ -- else aStar newAgenda (S.insert cl closed)
+ else aStar beamAgenda (S.insert cl closed)
+
+
+candidates :: SearchState s => Agendum s -> ExploredStates s -> CaveContext (Q.Seq (Agendum s))
+candidates agendum closed =
+ do let candidate = agendum ^. current
+ let previous = agendum ^. trail
+ let prevBenefit = agendum ^. trailBenefit
+ succs <- successors candidate
+ succAgs <- mapM (makeAgendum previous prevBenefit) succs
+ let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit, Q.length $ s ^. trail) `S.notMember` closed) succAgs
+ return nonloops
+
+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]
+
+makeAgendum :: SearchState s => Q.Seq s -> Int -> s -> CaveContext (Agendum s)
+makeAgendum previous prevBenefit newState =
+ do predicted <- estimateBenefit newState (Q.length previous)
+ cf <- currentFlow newState
+ let newTrail = previous |> newState
+ let incurred = prevBenefit + cf
+ return Agendum { _current = newState
+ , _trail = newTrail
+ , _trailBenefit = incurred
+ , _benefit = incurred + predicted
+ }
+
+
+isGoal :: SearchState s => Agendum s -> CaveContext Bool
+isGoal agendum =
+ do timeLimit <- asks getTimeLimit
+ return $ Q.length (agendum ^. trail) == (timeLimit - 1)
+
+isFullFlow :: SearchState s => s -> CaveContext Bool
+isFullFlow state =
+ do cave <- asks getCave
+ cf <- currentFlow state
+ let ff = sumOf (folded . flowRate) cave
+ return (cf == ff)
+
+
+-- Parse the input file
+
+caveP :: Parser Cave
+valveP :: Parser (RoomID, Room)
+roomP :: Parser Room
+tunnelsP :: Parser [RoomID]
+tunnelTextP :: Parser Text
+
+caveP = M.fromList <$> valveP `sepBy` endOfLine
+valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
+roomP = roomify <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP)
+ where roomify v ts = Room {_flowRate = v, _tunnels = ts }
+tunnelsP = (many1 letter) `sepBy` ", "
+tunnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
+
+successfulParse :: Text -> Cave
+successfulParse input =
+ case parseOnly caveP input of
+ Left _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
+ Right cave -> cave
\ No newline at end of file
projects / advent-of-code-22.git / blobdiff