advent16/MainCustomClosed.hs

   1 -- Writeup at https://work.njae.me.uk/2022/12/17/advent-of-code-2022-day-16/
   2
   3 import Debug.Trace
   4
   5 import AoC
   6 import Data.Text (Text)
   7 import qualified Data.Text.IO as TIO
   8 import Data.Attoparsec.Text hiding (take, D)
   9 import Control.Applicative
  10 import qualified Data.PQueue.Prio.Max as P
  11 import qualified Data.Set as S
  12 import qualified Data.Sequence as Q
  13 import qualified Data.Map.Strict as M
  14 import Data.Map.Strict ((!))
  15 -- import Data.Sequence ((|>), Seq((:|>)), ViewR ((:>)))
  16 import Data.Sequence ( (<|), (|>), Seq((:|>)) )
  17 import Data.List
  18 import Data.List.Split (chunksOf)
  19 import Data.Ord
  20 import Control.Monad.Reader
  21 import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
  22
  23
  24 type RoomID = String
  25
  26 data Tunnel = Tunnel { _tunnelTo :: RoomID, _tunnelLength :: Int}
  27     deriving (Eq, Show, Ord)
  28 makeLenses ''Tunnel
  29
  30 data Room = Room
  31     { _flowRate :: Int
  32     , _tunnels :: S.Set Tunnel
  33     } deriving (Eq, Show, Ord)
  34 makeLenses ''Room
  35
  36 type Cave = M.Map RoomID Room
  37 data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int , getSortedRooms :: [RoomID]}
  38
  39 type CaveContext = Reader TimedCave
  40
  41 data SingleSearchState = SingleSearchState
  42     { _currentRoom :: RoomID
  43     , _currentTime :: Int
  44     , _sOpenValves :: [RoomID]
  45     } deriving (Eq, Show, Ord)
  46 makeLenses ''SingleSearchState
  47
  48 data DoubleSearchState = DoubleSearchState
  49     { _personRoom :: RoomID
  50     , _personTime :: Int
  51     , _elephantRoom :: RoomID
  52     , _elephantTime :: Int
  53     , _dOpenValves :: [RoomID]
  54     } deriving (Eq, Show, Ord)
  55 makeLenses ''DoubleSearchState
  56
  57 data Agendum s =
  58     Agendum { _current :: s
  59             , _trail :: Q.Seq s
  60             , _trailBenefit :: Int
  61             , _benefit :: Int
  62             } deriving (Show, Eq, Ord)
  63 makeLenses ''Agendum
  64
  65 type Agenda s = P.MaxPQueue Int (Agendum s)
  66
  67 -- state, total flowed so far
  68 -- type ExploredStates s = S.Set (s, Int)
  69
  70
  71 class (Eq s, Ord s, Show s) => SearchState s where
  72   emptySearchState :: RoomID -> s
  73   currentFlow :: s -> CaveContext Int
  74   timeOf :: s -> Int
  75   successors :: s -> CaveContext (Q.Seq s)
  76   -- estimateBenefit :: s -> Int -> CaveContext Int
  77   estimateBenefit :: s -> CaveContext Int
  78
  79   data ExploredStateKey s
  80   -- type ExploredStates s
  81
  82   mkExploredKey :: s -> (ExploredStateKey s)
  83
  84 -- type ExploredStates s = M.Map (ExploredStateKey s) Int -- room/valves to time
  85 type ExploredStates s = S.Set ((ExploredStateKey s), Int) -- room & valves, trail benefit
  86
  87 instance SearchState SingleSearchState where
  88   emptySearchState startID = SingleSearchState
  89         { _currentRoom = startID
  90         , _currentTime = 0
  91         , _sOpenValves = []
  92         }
  93
  94   data ExploredStateKey SingleSearchState = SingleExploredStateKey RoomID [RoomID] -- current room and open valves
  95     deriving (Show, Eq, Ord)
  96
  97   mkExploredKey s = SingleExploredStateKey (s ^. currentRoom) (s ^. sOpenValves)
  98
  99   currentFlow state =
 100     do cave <- asks getCave
 101        let valves = state ^. sOpenValves
 102        let presentRooms = cave `M.restrictKeys` (S.fromList valves)
 103        -- let presentRooms = M.filter (\v -> v `elem` valves) cave
 104        return $ sumOf (folded . flowRate) presentRooms
 105
 106   timeOf state = state ^. currentTime
 107
 108   successors state =
 109     do isFF <- isFullFlow state
 110        -- cave <- asks getCave
 111        timeLimit <- asks getTimeLimit
 112        let here = state ^. currentRoom
 113        let opened = state ^. sOpenValves
 114        let now = state ^. currentTime
 115        succs <- agentSuccessor now opened now here
 116        let succStates = Q.fromList succs
 117        if isFF || (Q.null succStates)
 118        then return $ Q.singleton (state & currentTime .~ timeLimit)
 119        else return succStates
 120
 121   estimateBenefit here =
 122     do cave <- asks getCave
 123        timeLimit <- asks getTimeLimit
 124        let timeRemaining = timeLimit - (timeOf here)
 125        cf <- currentFlow here
 126        -- let closedValves = (cave `M.withoutKeys` (here ^. sOpenValves)) ^.. folded . flowRate
 127        -- let sortedClosedValves = sortOn Down closedValves
 128        sortedValves <- asks getSortedRooms
 129        let opened = here ^. sOpenValves
 130        let sortedClosedValves = [(cave ! v) ^. flowRate | v <- sortedValves, v `notElem` opened]
 131        let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
 132        return $ (cf * timeRemaining) + otherValveFlows
 133
 134
 135 instance SearchState DoubleSearchState where
 136   emptySearchState startID = DoubleSearchState
 137       { _personRoom = startID
 138       , _personTime = 0
 139       , _elephantRoom = startID
 140       , _elephantTime = 0
 141       , _dOpenValves = []
 142       }
 143
 144   data ExploredStateKey DoubleSearchState = DoubleExploredStateKey RoomID RoomID [RoomID] -- current room (person, elephant) and open valves
 145     -- deriving (Show)
 146     deriving (Show, Eq, Ord)
 147   -- type ExploredStates DoubleSearchState = M.Map (DoubleExploredStateKey DoubleSearchState) Int -- room/valves to time
 148
 149   mkExploredKey s = DoubleExploredStateKey minRoom maxRoom (s ^. dOpenValves)
 150     where minRoom = min (s ^. personRoom) (s ^. elephantRoom)
 151           maxRoom = max (s ^. personRoom) (s ^. elephantRoom)
 152
 153   currentFlow state =
 154     do cave <- asks getCave
 155        -- let valves = S.toList $ state ^. dOpenValves
 156        let valves = state ^. dOpenValves
 157        return $ sum $ fmap (\v -> (cave ! v) ^. flowRate) valves
 158        -- let presentRooms = cave `M.restrictKeys` valves
 159        -- return $ sumOf (folded . flowRate) presentRooms
 160
 161   timeOf state = min (state ^. personTime) (state ^. elephantTime)
 162
 163   successors state =
 164     do isFF <- isFullFlow state
 165        -- cave <- asks getCave
 166        timeLimit <- asks getTimeLimit
 167        let opened = state ^. dOpenValves
 168        let pNow = state ^. personTime
 169        let eNow = state ^. elephantTime
 170        let now = min pNow eNow
 171        let pHere = state ^. personRoom
 172        let eHere = state ^. elephantRoom
 173        pNexts <- agentSuccessor now opened pNow pHere
 174        eNexts <- agentSuccessor now opened eNow eHere
 175        let nexts =  [ state & personRoom .~ (p ^. currentRoom)
 176                             & personTime .~ (p ^. currentTime)
 177                             & elephantRoom .~ (e ^. currentRoom)
 178                             & elephantTime .~ (e ^. currentTime)
 179                             -- & dOpenValves %~ (S.union (p ^. sOpenValves) . S.union (e ^. sOpenValves))
 180                             & dOpenValves .~ (union (union opened (p ^. sOpenValves)) (e ^. sOpenValves))
 181                     | p <- pNexts
 182                     , e <- eNexts
 183                     ]
 184        let dedups = if pNow == eNow && pHere == eHere
 185                     then filter (\s -> (s ^. personRoom) < (s ^. elephantRoom)) nexts
 186                     else nexts
 187        -- let succStates = trace ("Succs: in " ++ (show state) ++ " out " ++ (show dedups)) (Q.fromList dedups)
 188        let succStates = Q.fromList dedups
 189        if isFF || (Q.null succStates)
 190        then return $ Q.singleton (state & personTime .~ timeLimit & elephantTime .~ timeLimit)
 191        else return succStates
 192
 193   estimateBenefit here =
 194     do cave <- asks getCave
 195        timeLimit <- asks getTimeLimit
 196        let timeRemaining = timeLimit - (timeOf here)
 197        cf <- currentFlow here
 198        -- let closedValves = (cave `M.withoutKeys` (here ^. dOpenValves)) ^.. folded . flowRate
 199        -- let sortedClosedValves = fmap sum $ chunksOf 2 $  {-# SCC estSort #-} sortOn Down closedValves
 200        -- let sortedClosedValves = fmap sum $ chunksOf 2 $ reverse $ sort closedValves -- no significant improvement
 201        sortedValves <- asks getSortedRooms
 202        let opened = here ^. dOpenValves
 203        let sortedClosedValves = fmap sum $ chunksOf 2 $ [(cave ! v) ^. flowRate | v <- sortedValves, v `notElem` opened]
 204        let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
 205        -- let otherValveFlows = timeRemaining * (sum closedValves) -- 8 minute runtime rather than 1:50
 206        return $ (cf * timeRemaining) + otherValveFlows
 207
 208 -- instance Eq (ExploredStateKey DoubleSearchState) where
 209 --     (DoubleExploredStateKey r1a r1b v1) == (DoubleExploredStateKey r2a r2b v2) =
 210 --       -- ((r1a == r2a && r1b == r2b) || (r1a == r2b && r1b == r2a)) && v1 == v2
 211 --       ((min r1a r1b), (max r1a r1b), v1) == ((min r2a r2b), (max r2a r2b), v2)
 212 --   -- data instance Ord DoubleExploredStateKey where
 213 -- instance Ord (ExploredStateKey DoubleSearchState) where
 214 --     (DoubleExploredStateKey r1a r1b v1) `compare` (DoubleExploredStateKey r2a r2b v2) =
 215 --       ((min r1a r1b), (max r1a r1b), v1) `compare` ((min r2a r2b), (max r2a r2b), v2)
 216
 217
 218 main :: IO ()
 219 main =
 220   do  dataFileName <- getDataFileName
 221       text <- TIO.readFile dataFileName
 222       let expandedCave = successfulParse text
 223       -- print cave
 224       -- print $ reachableFrom cave [Tunnel "AA" 0] S.empty []
 225       -- print $ compress cave
 226       let cave = compress expandedCave
 227       print $ part1 cave
 228       print $ part2 cave
 229
 230 -- part1 :: Cave -> Maybe (Agendum SingleSearchState)
 231 -- part1 cave = runReader (searchCave "AA") (TimedCave cave 30)
 232
 233 -- part2 :: Cave -> Maybe (Agendum DoubleSearchState)
 234 -- part2 cave = runReader (searchCave "AA") (TimedCave cave 26)
 235
 236 part1, part2 :: Cave -> Int
 237 -- part1 :: Cave -> Int
 238 part1 cave = maybe 0 _benefit result
 239 -- part1 cave = result
 240     where result = runReader (searchCave "AA") (TimedCave cave 30 sortedRooms) :: Maybe (Agendum SingleSearchState)
 241           sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
 242 part2 cave = maybe 0 _benefit result
 243 -- part2 cave = result
 244     where result = runReader (searchCave "AA") (TimedCave cave 26 sortedRooms) :: Maybe (Agendum DoubleSearchState)
 245           sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
 246           -- sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys cave
 247
 248 searchCave :: ((Ord (ExploredStateKey s)), (Show (ExploredStateKey s)), SearchState s) => String -> CaveContext (Maybe (Agendum s))
 249 searchCave startRoom =
 250     do agenda <- initAgenda startRoom
 251        aStar agenda S.empty
 252
 253 initAgenda :: ((Ord (ExploredStateKey s)), (Show (ExploredStateKey s)), SearchState s) => String -> CaveContext (Agenda s)
 254 initAgenda startID =
 255     do let startState = emptySearchState startID
 256        b <- estimateBenefit startState
 257        return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
 258
 259 aStar :: ((Ord (ExploredStateKey s)), (Show (ExploredStateKey s)), SearchState s) => Agenda s -> ExploredStates s -> CaveContext (Maybe (Agendum s))
 260 aStar agenda closed
 261        -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
 262     -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda) ++ " : trail " ++ (show $ _trail $ snd $ P.findMax agenda) ++ " : closed " ++ (show closed)) False = undefined
 263        -- | trace ("Peeping " ++ (show $ P.findMax agenda)) False = undefined
 264     | P.null agenda = return Nothing
 265     | otherwise =
 266         do  let (_, currentAgendum) = P.findMax agenda
 267             let reached = currentAgendum ^. current
 268             nexts <- candidates currentAgendum closed
 269             let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
 270             -- let beamAgenda = P.fromDescList $ P.take 10000 newAgenda -- agenda beam width
 271             -- let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
 272             -- let beamAgenda = P.fromDescList $ P.take 1000 newAgenda -- agenda beam width
 273             reachedGoal <- isGoal currentAgendum
 274             -- let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
 275             -- let cl = (reached, currentAgendum ^. trailBenefit)
 276             let cl = (mkExploredKey reached, currentAgendum ^. trailBenefit)
 277             if reachedGoal
 278             then return (Just currentAgendum)
 279             else if (cl `elem` closed)
 280                  then aStar (P.deleteMax agenda) closed
 281                  else aStar newAgenda (S.insert cl closed)
 282
 283
 284 candidates :: ((Ord (ExploredStateKey s)), SearchState s) => Agendum s -> ExploredStates s -> CaveContext (Q.Seq (Agendum s))
 285 candidates agendum closed =
 286     do  let candidate = agendum ^. current
 287         let previous = agendum ^. trail
 288         let prevBenefit = agendum ^. trailBenefit
 289         succs <- successors candidate
 290         succAgs <- mapM (makeAgendum previous prevBenefit) succs
 291         -- let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit, Q.length $ s ^. trail) `S.notMember` closed) succAgs
 292         -- let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit) `S.notMember` closed) succAgs
 293         let nonloops = Q.filter (\l -> ((mkExploredKey (l ^. current)), l ^. trailBenefit) `notElem` closed) succAgs
 294         return nonloops
 295
 296
 297 agentSuccessor :: Int -> [RoomID] -> Int -> RoomID -> CaveContext [SingleSearchState]
 298 agentSuccessor now opened aTime here
 299   | aTime /= now = return [SingleSearchState { _currentRoom = here, _currentTime = aTime, _sOpenValves = opened }]
 300   | otherwise =
 301       do  cave <- asks getCave
 302           timeLimit <- asks getTimeLimit
 303           -- let remaining = S.toList $ S.filter (\t -> (t ^. tunnelTo) `S.notMember` opened) ((cave ! here) ^. tunnels)
 304           let remaining = [ t
 305                           | t <- (S.toList ((cave ! here) ^. tunnels))
 306                           , (t ^. tunnelTo) `notElem` opened
 307                           ]
 308           let moves = [ SingleSearchState
 309                             { _currentRoom = (t ^. tunnelTo)
 310                             , _currentTime = now + (t ^. tunnelLength)
 311                             , _sOpenValves = opened
 312                             }
 313                       | t <- remaining
 314                       , now + (t ^. tunnelLength) <= timeLimit
 315                       ]
 316           let moves' = ( SingleSearchState
 317                                 { _currentRoom = here
 318                                 , _currentTime = timeLimit
 319                                 , _sOpenValves = opened
 320                                 }
 321                        : moves)
 322           let opens = if here `notElem` opened && (cave ! here) ^. flowRate > 0
 323                       then [SingleSearchState { _currentRoom = here, _currentTime = aTime + 1, _sOpenValves = opened ++ [here] }]
 324                       else []
 325           -- let nexts = moves ++ opens
 326           let nexts = if null opens then moves' else opens
 327           return nexts
 328
 329 makeAgendum :: SearchState s => Q.Seq s -> Int -> s -> CaveContext (Agendum s)
 330 makeAgendum previous prevBenefit newState =
 331     do predicted <- estimateBenefit newState --  (Q.length previous)
 332        -- cf <- currentFlow newState
 333        oldFlow <- lastFlow previous (timeOf newState)
 334        let newTrail = previous |> newState
 335        let incurred = prevBenefit + oldFlow
 336        return Agendum { _current = newState
 337                       , _trail = newTrail
 338                       , _trailBenefit = incurred
 339                       , _benefit = incurred + predicted
 340                       }
 341
 342 lastFlow :: SearchState s => Q.Seq s -> Int -> CaveContext Int
 343 lastFlow Q.Empty _ = return 0
 344 lastFlow (_ :|> previous) newTime =
 345   do cf <- currentFlow previous
 346      let dt = newTime - (timeOf previous)
 347      return (cf * dt)
 348
 349 isGoal :: SearchState s => Agendum s -> CaveContext Bool
 350 isGoal agendum =
 351   do timeLimit <- asks getTimeLimit
 352      let s = agendum ^. current
 353      return $ (timeOf s) == timeLimit
 354
 355 isFullFlow :: SearchState s => s -> CaveContext Bool
 356 isFullFlow state =
 357   do cave <- asks getCave
 358      cf <- currentFlow state
 359      let ff = sumOf (folded . flowRate) cave
 360      return (cf == ff)
 361
 362 compress :: Cave -> Cave
 363 compress cave = M.mapWithKey (compressRoom cave) cave
 364
 365 compressRoom :: Cave -> RoomID -> Room -> Room
 366 compressRoom cave here room = room & tunnels .~ t'
 367   where t' = reachableFrom cave [Tunnel here 0] S.empty S.empty
 368
 369 reachableFrom :: Cave -> [Tunnel] -> S.Set RoomID -> S.Set Tunnel -> S.Set Tunnel
 370 reachableFrom _ [] _ routes = routes
 371 reachableFrom cave (tunnel@(Tunnel here len):boundary) found routes
 372   | here `S.member` found = reachableFrom cave boundary found routes
 373   | otherwise = reachableFrom cave (boundary ++ (S.toList legs)) (S.insert here found) routes'
 374   where exits = (cave ! here) ^. tunnels
 375         exits' = S.filter (\t -> (t ^. tunnelTo) `S.notMember` found) exits
 376         legs = S.map (\t -> t & tunnelLength .~ (len + 1)) exits'
 377         routes' = if (len == 0) || ((cave ! here) ^. flowRate) == 0
 378                   then routes
 379                   else S.insert tunnel routes
 380
 381
 382 -- Parse the input file
 383
 384 caveP :: Parser Cave
 385 valveP :: Parser (RoomID, Room)
 386 roomP :: Parser Room
 387 tunnelsP :: Parser (S.Set Tunnel)
 388 tunnelTextP :: Parser Text
 389
 390 caveP = M.fromList <$> valveP `sepBy` endOfLine
 391 valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
 392 roomP = Room <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP)
 393     -- where roomify v ts = Room {flowRate = v, tunnels = ts }
 394 tunnelsP = (S.fromList . (fmap (flip Tunnel 1))) <$> (many1 letter) `sepBy` ", "
 395 tunnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
 396
 397 successfulParse :: Text -> Cave
 398 successfulParse input =
 399   case parseOnly caveP input of
 400     Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
 401     Right cave -> cave