advent16/Main.hs

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