X-Git-Url: https://git.njae.me.uk/?p=advent-of-code-22.git;a=blobdiff_plain;f=advent16%2FMainEstSort.hs;fp=advent16%2FMainEstSort.hs;h=459a10be5fe02a723ddf459125206a8c5cfebefa;hp=0000000000000000000000000000000000000000;hb=4087698696ed09477c3b5073f3d4d93d85c0a632;hpb=549425defbc1482abcef0e926094f0817842a4f5 diff --git a/advent16/MainEstSort.hs b/advent16/MainEstSort.hs new file mode 100644 index 0000000..459a10b --- /dev/null +++ b/advent16/MainEstSort.hs @@ -0,0 +1,352 @@ +-- 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 ((|>), Seq((:|>)), ViewR ((:>))) +import Data.Sequence ( (|>), Seq((:|>)) ) +import Data.List +import Data.List.Split (chunksOf) +import Data.Ord +import Control.Monad.Reader +import Control.Lens hiding ((<|), (|>), (:>), (:<), indices) + + +type RoomID = String + +data Tunnel = Tunnel { _tunnelTo :: RoomID, _tunnelLength :: Int} + deriving (Eq, Show, Ord) +makeLenses ''Tunnel + +data Room = Room + { _flowRate :: Int + , _tunnels :: S.Set Tunnel + } 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 + , _currentTime :: Int + , _sOpenValves :: S.Set RoomID + } deriving (Eq, Show, Ord) +makeLenses ''SingleSearchState + +data DoubleSearchState = DoubleSearchState + { _personRoom :: RoomID + , _personTime :: Int + , _elephantRoom :: RoomID + , _elephantTime :: Int + , _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) + +-- state, total flowed so far +type ExploredStates s = S.Set (s, Int) + + +class (Eq s, Ord s, Show s) => SearchState s where + emptySearchState :: RoomID -> s + currentFlow :: s -> CaveContext Int + timeOf :: s -> Int + successors :: s -> CaveContext (Q.Seq s) + -- estimateBenefit :: s -> Int -> CaveContext Int + estimateBenefit :: s -> CaveContext Int + +instance SearchState SingleSearchState where + emptySearchState startID = SingleSearchState + { _currentRoom = startID + , _currentTime = 0 + , _sOpenValves = S.empty + } + + currentFlow state = + do cave <- asks getCave + let valves = state ^. sOpenValves + let presentRooms = cave `M.restrictKeys` valves + return $ sumOf (folded . flowRate) presentRooms + + timeOf state = state ^. currentTime + + successors state = + do isFF <- isFullFlow state + -- cave <- asks getCave + timeLimit <- asks getTimeLimit + let here = state ^. currentRoom + let opened = state ^. sOpenValves + let now = state ^. currentTime + succs <- agentSuccessor now opened now here + let succStates = Q.fromList succs + if isFF || (Q.null succStates) + then return $ Q.singleton (state & currentTime .~ timeLimit) + else return succStates + + estimateBenefit here = + do cave <- asks getCave + timeLimit <- asks getTimeLimit + let timeRemaining = timeLimit - (timeOf here) + 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 + , _personTime = 0 + , _elephantRoom = startID + , _elephantTime = 0 + , _dOpenValves = S.empty + } + + currentFlow state = + do cave <- asks getCave + let valves = S.toList $ state ^. dOpenValves + return $ sum $ fmap (\v -> (cave ! v) ^. flowRate) valves + -- let presentRooms = cave `M.restrictKeys` valves + -- return $ sumOf (folded . flowRate) presentRooms + + timeOf state = min (state ^. personTime) (state ^. elephantTime) + + successors state = + do isFF <- isFullFlow state + -- cave <- asks getCave + timeLimit <- asks getTimeLimit + let opened = state ^. dOpenValves + let pNow = state ^. personTime + let eNow = state ^. elephantTime + let now = min pNow eNow + let pHere = state ^. personRoom + let eHere = state ^. elephantRoom + pNexts <- agentSuccessor now opened pNow pHere + eNexts <- agentSuccessor now opened eNow eHere + let nexts = [ state & personRoom .~ (p ^. currentRoom) + & personTime .~ (p ^. currentTime) + & elephantRoom .~ (e ^. currentRoom) + & elephantTime .~ (e ^. currentTime) + & dOpenValves %~ (S.union (p ^. sOpenValves) . S.union (e ^. sOpenValves)) + | p <- pNexts + , e <- eNexts + ] + let dedups = if pNow == eNow && pHere == eHere + then filter (\s -> (s ^. personRoom) < (s ^. elephantRoom)) nexts + else nexts + -- let succStates = trace ("Succs: in " ++ (show state) ++ " out " ++ (show dedups)) (Q.fromList dedups) + let succStates = Q.fromList dedups + if isFF || (Q.null succStates) + then return $ Q.singleton (state & personTime .~ timeLimit & elephantTime .~ timeLimit) + else return succStates + + estimateBenefit here = + do cave <- asks getCave + timeLimit <- asks getTimeLimit + let timeRemaining = timeLimit - (timeOf here) + cf <- currentFlow here + let closedValves = (cave `M.withoutKeys` (here ^. dOpenValves)) ^.. folded . flowRate + let sortedClosedValves = fmap sum $ chunksOf 2 $ {-# SCC estSort #-} sortOn Down closedValves + -- let sortedClosedValves = fmap sum $ chunksOf 2 $ reverse $ sort closedValves -- no significant improvement + let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves + -- let otherValveFlows = timeRemaining * (sum closedValves) -- 8 minute runtime rather than 1:50 + return $ (cf * timeRemaining) + otherValveFlows + + +main :: IO () +main = + do dataFileName <- getDataFileName + text <- TIO.readFile dataFileName + let expandedCave = successfulParse text + -- print cave + -- print $ reachableFrom cave [Tunnel "AA" 0] S.empty [] + -- print $ compress cave + let cave = compress expandedCave + 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 -> 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 + 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) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda)) False = undefined + -- | trace ("Peeping " ++ (show $ 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 + -- let beamAgenda = P.fromDescList $ P.take 1000 newAgenda -- agenda beam width + reachedGoal <- isGoal currentAgendum + -- let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail) + let cl = (reached, currentAgendum ^. trailBenefit) + 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 + let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit) `S.notMember` closed) succAgs + return nonloops + + +agentSuccessor :: Int -> S.Set RoomID -> Int -> RoomID -> CaveContext [SingleSearchState] +agentSuccessor now opened aTime here + | aTime /= now = return [SingleSearchState { _currentRoom = here, _currentTime = aTime, _sOpenValves = opened }] + | otherwise = + do cave <- asks getCave + timeLimit <- asks getTimeLimit + let remaining = S.toList $ S.filter (\t -> (t ^. tunnelTo) `S.notMember` opened) ((cave ! here) ^. tunnels) + let moves = [ SingleSearchState + { _currentRoom = (t ^. tunnelTo) + , _currentTime = now + (t ^. tunnelLength) + , _sOpenValves = opened + } + | t <- remaining + , now + (t ^. tunnelLength) <= timeLimit + ] + let opens = if here `S.notMember` opened && (cave ! here) ^. flowRate > 0 + then [SingleSearchState { _currentRoom = here, _currentTime = aTime + 1, _sOpenValves = S.insert here opened }] + else [] + -- let nexts = moves ++ opens + let nexts = if null opens then moves else opens + let nexts' = if null nexts + then [ SingleSearchState + { _currentRoom = here + , _currentTime = timeLimit + , _sOpenValves = opened + } ] + else nexts + return nexts' + +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 + oldFlow <- lastFlow previous (timeOf newState) + let newTrail = previous |> newState + let incurred = prevBenefit + oldFlow + return Agendum { _current = newState + , _trail = newTrail + , _trailBenefit = incurred + , _benefit = incurred + predicted + } + +lastFlow :: SearchState s => Q.Seq s -> Int -> CaveContext Int +lastFlow Q.Empty _ = return 0 +lastFlow (_ :|> previous) newTime = + do cf <- currentFlow previous + let dt = newTime - (timeOf previous) + return (cf * dt) + +isGoal :: SearchState s => Agendum s -> CaveContext Bool +isGoal agendum = + do timeLimit <- asks getTimeLimit + let s = agendum ^. current + return $ (timeOf s) == timeLimit + +isFullFlow :: SearchState s => s -> CaveContext Bool +isFullFlow state = + do cave <- asks getCave + cf <- currentFlow state + let ff = sumOf (folded . flowRate) cave + return (cf == ff) + +compress :: Cave -> Cave +compress cave = M.mapWithKey (compressRoom cave) cave + +compressRoom :: Cave -> RoomID -> Room -> Room +compressRoom cave here room = room & tunnels .~ t' + where t' = reachableFrom cave [Tunnel here 0] S.empty S.empty + +reachableFrom :: Cave -> [Tunnel] -> S.Set RoomID -> S.Set Tunnel -> S.Set Tunnel +reachableFrom _ [] _ routes = routes +reachableFrom cave (tunnel@(Tunnel here len):boundary) found routes + | here `S.member` found = reachableFrom cave boundary found routes + | otherwise = reachableFrom cave (boundary ++ (S.toList legs)) (S.insert here found) routes' + where exits = (cave ! here) ^. tunnels + exits' = S.filter (\t -> (t ^. tunnelTo) `S.notMember` found) exits + legs = S.map (\t -> t & tunnelLength .~ (len + 1)) exits' + routes' = if (len == 0) || ((cave ! here) ^. flowRate) == 0 + then routes + else S.insert tunnel routes + + +-- Parse the input file + +caveP :: Parser Cave +valveP :: Parser (RoomID, Room) +roomP :: Parser Room +tunnelsP :: Parser (S.Set Tunnel) +tunnelTextP :: Parser Text + +caveP = M.fromList <$> valveP `sepBy` endOfLine +valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP +roomP = Room <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP) + -- where roomify v ts = Room {flowRate = v, tunnels = ts } +tunnelsP = (S.fromList . (fmap (flip Tunnel 1))) <$> (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