advent16/Main.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 ((|>))
  16 import Data.List
  17 import Data.List.Split (chunksOf)
  18 import Data.Ord
  19 import Control.Monad.Reader
  20 import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
  21
  22 -- pattern Empty   <- (Q.viewl -> Q.EmptyL)  where Empty = Q.empty
  23 -- pattern x :< xs <- (Q.viewl -> x Q.:< xs) where (:<)  = (Q.<|)
  24 -- pattern xs :> x <- (Q.viewr -> xs Q.:> x) where (:>)  = (Q.|>)
  25
  26 type RoomID = String
  27
  28 data Room = Room
  29     { _flowRate :: Int
  30     , _tunnels :: [RoomID]
  31     } deriving (Eq, Show, Ord)
  32 makeLenses ''Room
  33
  34 type Cave = M.Map RoomID Room
  35 data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int}
  36
  37 type CaveContext = Reader TimedCave
  38
  39 data SingleSearchState = SingleSearchState
  40     { _currentRoom :: RoomID
  41     , _sOpenValves :: S.Set RoomID
  42     } deriving (Eq, Show, Ord)
  43 makeLenses ''SingleSearchState
  44
  45 data DoubleSearchState = DoubleSearchState
  46     { _personRoom :: RoomID
  47     , _elephantRoom :: RoomID
  48     , _dOpenValves :: S.Set RoomID
  49     } deriving (Eq, Show, Ord)
  50 makeLenses ''DoubleSearchState
  51
  52 data Agendum s =
  53     Agendum { _current :: s
  54             , _trail :: Q.Seq s
  55             , _trailBenefit :: Int
  56             , _benefit :: Int
  57             } deriving (Show, Eq, Ord)
  58 makeLenses ''Agendum
  59
  60 type Agenda s = P.MaxPQueue Int (Agendum s)
  61
  62 type ExploredStates s = S.Set (s, Int, Int)
  63
  64
  65 class (Eq s, Ord s, Show s) => SearchState s where
  66   emptySearchState :: RoomID -> s
  67   currentFlow :: s -> CaveContext Int
  68   successors :: s -> CaveContext (Q.Seq s)
  69   estimateBenefit :: s -> Int -> CaveContext Int
  70
  71 instance SearchState SingleSearchState where
  72   emptySearchState startID = SingleSearchState { _currentRoom = startID, _sOpenValves = S.empty }
  73
  74   currentFlow state =
  75     do cave <- asks getCave
  76        let valves = state ^. sOpenValves
  77        let presentRooms = cave `M.restrictKeys` valves
  78        return $ sumOf (folded . flowRate) presentRooms
  79
  80   successors state =
  81     do isFF <- isFullFlow state
  82        let here = state ^. currentRoom
  83        let opened = state ^. sOpenValves
  84        succPairs <- personSuccessor here opened
  85        let succStates =
  86               [ SingleSearchState
  87                   { _currentRoom = r
  88                   , _sOpenValves = o
  89                   }
  90               | (r, o) <- succPairs
  91               ]
  92        if isFF
  93        then return $ Q.singleton state
  94        else return $ Q.fromList succStates
  95
  96   estimateBenefit here timeElapsed =
  97     do cave <- asks getCave
  98        timeLimit <- asks getTimeLimit
  99        let timeRemaining = timeLimit - (timeElapsed + 2)
 100        cf <- currentFlow here
 101        let closedValves = (cave `M.withoutKeys` (here ^. sOpenValves)) ^.. folded . flowRate
 102        let sortedClosedValves = sortOn Down closedValves
 103        let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
 104        return $ (cf * timeRemaining) + otherValveFlows
 105
 106
 107 instance SearchState DoubleSearchState where
 108   emptySearchState startID = DoubleSearchState
 109       { _personRoom = startID
 110       , _elephantRoom = startID
 111       , _dOpenValves = S.empty
 112       }
 113
 114   currentFlow state =
 115     do cave <- asks getCave
 116        let valves = state ^. dOpenValves
 117        let presentRooms = cave `M.restrictKeys` valves
 118        return $ sumOf (folded . flowRate) presentRooms
 119
 120   successors state =
 121     do isFF <- isFullFlow state
 122        let pHere = state ^. personRoom
 123        let eHere = state ^. elephantRoom
 124        let opened = state ^. dOpenValves
 125        pSuccPairs <- personSuccessor pHere opened
 126        eSuccPairs <- personSuccessor eHere opened
 127        let succStates =
 128               [ DoubleSearchState
 129                   { _personRoom = p
 130                   , _elephantRoom = e
 131                   , _dOpenValves = S.union po eo
 132                   }
 133               | (p, po) <- pSuccPairs
 134               , (e, eo) <- eSuccPairs
 135               ]
 136        if isFF
 137        then return $ Q.singleton state
 138        else return $ Q.fromList succStates
 139
 140   estimateBenefit here timeElapsed =
 141     do cave <- asks getCave
 142        timeLimit <- asks getTimeLimit
 143        let timeRemaining = timeLimit - (timeElapsed + 2)
 144        cf <- currentFlow here
 145        let closedValves = (cave `M.withoutKeys` (here ^. dOpenValves)) ^.. folded . flowRate
 146        let sortedClosedValves = fmap sum $ chunksOf 2 $ sortOn Down closedValves
 147        let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
 148        return $ (cf * timeRemaining) + otherValveFlows
 149
 150
 151 main :: IO ()
 152 main =
 153   do  dataFileName <- getDataFileName
 154       text <- TIO.readFile dataFileName
 155       let cave = successfulParse text
 156       -- print cave
 157       print $ part1 cave
 158       print $ part2 cave
 159
 160 -- part1 :: Cave -> Maybe (Agendum SingleSearchState)
 161 -- part1 cave = runReader (searchCave "AA") (TimedCave cave 30)
 162
 163 -- part2 :: Cave -> Maybe (Agendum DoubleSearchState)
 164 -- part2 cave = runReader (searchCave "AA") (TimedCave cave 26)
 165
 166 part1, part2 :: Cave -> Int
 167 part1 cave = maybe 0 _benefit result
 168     where result = runReader (searchCave "AA") (TimedCave cave 30) :: Maybe (Agendum SingleSearchState)
 169 part2 cave = maybe 0 _benefit result
 170     where result = runReader (searchCave "AA") (TimedCave cave 26) :: Maybe (Agendum DoubleSearchState)
 171
 172 searchCave :: SearchState s => String -> CaveContext (Maybe (Agendum s))
 173 searchCave startRoom =
 174     do agenda <- initAgenda startRoom
 175        aStar agenda S.empty
 176
 177 initAgenda :: SearchState s => String -> CaveContext (Agenda s)
 178 initAgenda startID =
 179     do let startState = emptySearchState startID
 180        b <- estimateBenefit startState 0
 181        return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
 182
 183 aStar ::  SearchState s => Agenda s -> ExploredStates s -> CaveContext (Maybe (Agendum s))
 184 aStar agenda closed
 185     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
 186     -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : len " ++ (show $ Q.length $ _trail $ snd $ P.findMax agenda)) False = undefined
 187     | P.null agenda = return Nothing
 188     | otherwise =
 189         do  let (_, currentAgendum) = P.findMax agenda
 190             let reached = currentAgendum ^. current
 191             nexts <- candidates currentAgendum closed
 192             let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
 193             -- let beamAgenda = P.fromDescList $ P.take 10000 newAgenda -- agenda beam width
 194             let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
 195             reachedGoal <- isGoal currentAgendum
 196             let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
 197             if reachedGoal
 198             then return (Just currentAgendum)
 199             else if (cl `S.member` closed)
 200                  then aStar (P.deleteMax agenda) closed
 201                  -- else aStar newAgenda (S.insert cl closed)
 202                  else aStar beamAgenda (S.insert cl closed)
 203
 204
 205 candidates :: SearchState s => Agendum s -> ExploredStates s -> CaveContext (Q.Seq (Agendum s))
 206 candidates agendum closed =
 207     do  let candidate = agendum ^. current
 208         let previous = agendum ^. trail
 209         let prevBenefit = agendum ^. trailBenefit
 210         succs <- successors candidate
 211         succAgs <- mapM (makeAgendum previous prevBenefit) succs
 212         let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit, Q.length $ s ^. trail) `S.notMember` closed) succAgs
 213         return nonloops
 214
 215 personSuccessor, openValveSuccessor, walkSuccessor ::  RoomID -> S.Set RoomID -> CaveContext [(RoomID, S.Set RoomID)]
 216 personSuccessor here opened =
 217   do ovs <- openValveSuccessor here opened
 218      ws <- walkSuccessor here opened
 219      return (ovs ++ ws)
 220
 221 openValveSuccessor here opened
 222   | here `S.member` opened = return []
 223   | otherwise = return [(here, S.insert here opened)]
 224
 225 walkSuccessor here opened =
 226   do cave <- asks getCave
 227      let neighbours = (cave ! here) ^. tunnels
 228      return [(n, opened) | n <- neighbours]
 229
 230 makeAgendum :: SearchState s => Q.Seq s -> Int -> s -> CaveContext (Agendum s)
 231 makeAgendum previous prevBenefit newState =
 232     do predicted <- estimateBenefit newState (Q.length previous)
 233        cf <- currentFlow newState
 234        let newTrail = previous |> newState
 235        let incurred = prevBenefit + cf
 236        return Agendum { _current = newState
 237                       , _trail = newTrail
 238                       , _trailBenefit = incurred
 239                       , _benefit = incurred + predicted
 240                       }
 241
 242
 243 isGoal :: SearchState s => Agendum s -> CaveContext Bool
 244 isGoal agendum =
 245   do timeLimit <- asks getTimeLimit
 246      return $ Q.length (agendum ^. trail) == (timeLimit - 1)
 247
 248 isFullFlow :: SearchState s => s -> CaveContext Bool
 249 isFullFlow state =
 250   do cave <- asks getCave
 251      cf <- currentFlow state
 252      let ff = sumOf (folded . flowRate) cave
 253      return (cf == ff)
 254
 255
 256 -- Parse the input file
 257
 258 caveP :: Parser Cave
 259 valveP :: Parser (RoomID, Room)
 260 roomP :: Parser Room
 261 tunnelsP :: Parser [RoomID]
 262 turnnelTextP :: Parser Text
 263
 264 caveP = M.fromList <$> valveP `sepBy` endOfLine
 265 valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
 266 roomP = roomify <$> (" has flow rate=" *> decimal) <*> (turnnelTextP *> tunnelsP)
 267     where roomify v ts = Room {_flowRate = v, _tunnels = ts }
 268 tunnelsP = (many1 letter) `sepBy` ", "
 269 turnnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
 270
 271 successfulParse :: Text -> Cave
 272 successfulParse input =
 273   case parseOnly caveP input of
 274     Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
 275     Right cave -> cave