advent18/src/advent18class.hs

   1 import Debug.Trace
   2
   3 -- import qualified Data.Text.IO as TIO
   4
   5 import qualified Data.Map.Strict as M
   6 import Data.Map.Strict ((!))
   7 import qualified Data.PQueue.Prio.Min as P
   8 import qualified Data.Set as S
   9 import qualified Data.Sequence as Q
  10 import Data.Sequence ((<|), (|>), (><))
  11 import Data.Foldable (toList, foldr', foldl', all)
  12 -- import Data.Maybe (fromJust)
  13 -- import Data.List
  14 import Data.Char
  15 import Control.Monad.Reader
  16 import Control.Lens hiding ((<|), (|>))
  17 -- import Data.Map.Lens
  18
  19
  20 type Position = (Integer, Integer) -- r, c
  21
  22 type Keys = S.Set Char
  23 type PointOfInterest = M.Map Position Char
  24
  25 data Explorer1 = Explorer1 { _explorer1Position :: Char
  26                            , _explorer1KeysHeld :: Keys
  27                            , _explorer1Travelled :: Int
  28                            } deriving (Show)
  29 data Explorer4 = Explorer4 { _explorer4Position :: S.Set Char
  30                            , _explorer4KeysHeld :: Keys
  31                            , _explorer4Travelled :: Int
  32                            } deriving (Show)
  33 makeFields ''Explorer1
  34 makeFields ''Explorer4
  35
  36 type ExploredStates e = S.Set e
  37
  38 type ExpandedCave = S.Set Position
  39 data ExpandedCaveComplex = ExpandedCaveComplex { _caveE :: ExpandedCave
  40                                , _keysE :: PointOfInterest
  41                                , _doors :: PointOfInterest
  42                                } deriving (Eq, Ord, Show)
  43 makeLenses ''ExpandedCaveComplex
  44
  45 data CaveEdge = CaveEdge { _keysRequired :: S.Set Char
  46                          , _distance :: Int
  47                          } deriving (Eq, Ord, Show)
  48 makeLenses ''CaveEdge
  49
  50 type EdgeKey = (Char, Char)
  51 type Cave = M.Map EdgeKey CaveEdge
  52
  53 data CaveComplex = CaveComplex { _cave :: Cave
  54                                , _keys :: S.Set Char
  55                                } deriving (Eq, Ord, Show)
  56 makeLenses ''CaveComplex
  57
  58 type CaveContext = Reader CaveComplex
  59
  60 data Agendum e = Agendum { _current :: e
  61                        , _trail :: Q.Seq e
  62                        , _cost :: Int} deriving (Show, Eq)
  63 type Agenda e = P.MinPQueue Int (Agendum e)
  64
  65
  66 instance Eq Explorer1 where
  67     e1 == e2 = (_explorer1Position e1 == _explorer1Position e2) && (_explorer1KeysHeld e1 == _explorer1KeysHeld e2)
  68 instance Eq Explorer4 where
  69     e1 == e2 = (_explorer4Position e1 == _explorer4Position e2) && (_explorer4KeysHeld e1 == _explorer4KeysHeld e2)
  70
  71 instance Ord Explorer1 where
  72     e1 `compare` e2 =
  73         if _explorer1Position e1 == _explorer1Position e2
  74         then (_explorer1KeysHeld e1) `compare` (_explorer1KeysHeld e2)
  75         else (_explorer1Position e1) `compare`(_explorer1Position e2)
  76 instance Ord Explorer4 where
  77     e1 `compare` e2 =
  78         if _explorer4Position e1 == _explorer4Position e2
  79         then (_explorer4KeysHeld e1) `compare` (_explorer4KeysHeld e2)
  80         else (_explorer4Position e1) `compare`(_explorer4Position e2)
  81
  82
  83 class (Eq e, Ord e, Show e) => Explorer e where
  84     successors :: e -> CaveContext (Q.Seq e)
  85     estimateCost :: e -> CaveContext Int
  86     extendExplorer :: e -> EdgeKey -> CaveEdge -> e
  87     -- positionE :: e -> Position
  88     -- keysHeldE :: e -> Keys
  89     emptyExplorer :: e
  90
  91 instance Explorer Explorer1 where
  92     successors explorer = -- return Q.empty
  93         do  let here = explorer ^. position
  94             cavern <- asks _cave
  95             let kH = explorer ^. keysHeld
  96             let locations0 = M.filterWithKey (\k _ds -> edgeTouches here k) cavern
  97             let locations1 = M.filter (\e -> S.null ((e ^. keysRequired) `S.difference` kH)) locations0
  98             let succs = M.foldrWithKey' (\k e q -> (extendExplorer explorer k e) <| q) Q.empty locations1
  99             return succs
 100
 101     estimateCost explorer = -- return 0
 102         do let here = explorer ^. position
 103            ks <- asks _keys
 104            cavern <- asks _cave
 105            let kH = explorer ^. keysHeld
 106            let unfound = ks `S.difference` kH
 107            let unfoundEdges = M.filterWithKey (\k _ -> (edgeTouches here k) && ((edgeOther here k) `S.member` unfound)) cavern
 108            let furthest = maximum $ (0:) $ map _distance $ M.elems unfoundEdges
 109            return $ max 0 $ furthest + (S.size unfound) - 1
 110            -- return $ S.size unfound
 111
 112     emptyExplorer = Explorer1 { _explorer1Position = '0', _explorer1KeysHeld = S.empty, _explorer1Travelled = 0 }
 113
 114     extendExplorer explorer edgeKey edge =
 115         explorer & position .~ there
 116                  & keysHeld .~ kH'
 117                  & travelled .~ d'
 118         where there = edgeOther (explorer ^. position) edgeKey
 119               kH' = S.insert there (explorer ^. keysHeld)
 120               d' = (explorer ^. travelled) + (edge ^. distance)
 121
 122 instance Explorer Explorer4 where
 123     successors explorer = -- return Q.empty
 124         do let heres = explorer ^. position
 125            cavern <- asks _cave
 126            let kH = explorer ^. keysHeld
 127            let locations0 = M.filterWithKey (\k _ds -> anyEdgeTouch heres k) cavern
 128            let locations1 = M.filter (\e -> S.null ((e ^. keysRequired) `S.difference` kH)) locations0
 129            let succs = M.foldrWithKey' (\k e q -> (extendExplorer explorer k e) <| q) Q.empty locations1
 130            return succs
 131
 132     estimateCost explorer = -- return 0
 133         do let heres = explorer ^. position
 134            ks <- asks _keys
 135            cavern <- asks _cave
 136            let kH = explorer ^. keysHeld
 137            let unfound = ks `S.difference` kH
 138            let unfoundEdges0 = M.filterWithKey (\k _ -> anyEdgeTouch heres k) cavern
 139            let unfoundEdges = M.filterWithKey (\k _ -> not $ anyEdgeTouch kH k) unfoundEdges0
 140            let furthest = maximum $ (0:) $ map _distance $ M.elems unfoundEdges
 141            return $ max 0 $ furthest + (S.size unfound) - 1
 142
 143     emptyExplorer = Explorer4 { _explorer4Position = S.fromList "0123", _explorer4KeysHeld = S.empty, _explorer4Travelled = 0 }
 144
 145     extendExplorer explorer edgeKey edge =
 146         explorer & position .~ pos'
 147                  & keysHeld .~ kH'
 148                  & travelled .~ d'
 149         where here = S.findMin $ S.filter (\p -> edgeTouches p edgeKey) (explorer ^. position)
 150               there = edgeOther here edgeKey
 151               kH' = S.insert there (explorer ^. keysHeld)
 152               d' = (explorer ^. travelled) + (edge ^. distance)
 153               pos' = S.insert there $ S.delete here (explorer ^. position)
 154
 155
 156 main :: IO ()
 157 main = do
 158         text <- readFile "data/advent18.txt"
 159         let (ccE, startPosition) = buildCaveComplex text
 160         -- print ccE
 161         print $ part1 ccE startPosition
 162         print $ part2 ccE startPosition
 163
 164
 165 part1 :: ExpandedCaveComplex -> Position -> Int
 166 part1 cavern startPosition = maybe 0 _cost result
 167     where cc = contractCave cavern [startPosition]
 168           explorer = emptyExplorer :: Explorer1
 169           result = runReader (searchCave explorer) cc
 170
 171 part2 ::  ExpandedCaveComplex -> Position -> Int
 172 part2 caveComplex0 (re, ce) = maybe 0 _cost result
 173     where
 174         startPositions = [(re - 1, ce - 1), (re - 1, ce + 1), (re + 1 , ce - 1), (re + 1, ce + 1)]
 175         cavern0 = _caveE caveComplex0
 176         cavern = cavern0 `S.difference` [(re, ce), (re + 1, ce), (re - 1, ce), (re, ce + 1), (re, ce - 1)]
 177         caveComplex = caveComplex0 {_caveE = cavern}
 178         cc = contractCave caveComplex startPositions
 179         explorer = emptyExplorer :: Explorer4
 180         result = runReader (searchCave explorer) cc
 181
 182
 183 -- buildCaveComplex :: Explorer e => String -> (CaveComplex, e)
 184 buildCaveComplex :: String -> (ExpandedCaveComplex, Position)
 185 buildCaveComplex text = (ccE, startPosition)
 186     where (ccE, startPosition) = foldl' buildCaveRow (cc0, (0, 0)) $ zip [0..] rows
 187           cc0 = ExpandedCaveComplex {_caveE = S.empty, _keysE = M.empty, _doors = M.empty}
 188           -- explorer0 = emptyExplorer -- Explorer { _position = (0, 0), _keysHeld = S.empty }
 189           rows = lines text
 190
 191 buildCaveRow :: (ExpandedCaveComplex, Position) -> (Integer, String) -> (ExpandedCaveComplex, Position)
 192 buildCaveRow (cc, explorers) (r, row) = foldl' (buildCaveCell r) (cc, explorers) $ zip [0..] row
 193
 194
 195 buildCaveCell :: Integer -> (ExpandedCaveComplex, Position) -> (Integer, Char) -> (ExpandedCaveComplex, Position)
 196 buildCaveCell r (cc, startPosition) (c, char)
 197     | char == '.' = (cc', startPosition)
 198     | char == '@' = (cc', here)
 199     | isLower char = (cc' { _keysE = M.insert here char $ _keysE cc'}, startPosition)
 200     | isUpper char = (cc' { _doors = M.insert here char $ _doors cc'}, startPosition)
 201     | otherwise = (cc, startPosition)
 202     where cc' = cc { _caveE = S.insert here $ _caveE cc }
 203           here = (r, c)
 204
 205
 206
 207 mkEdgeKey a b = if a < b then (a, b) else (b, a)
 208
 209 edgeTouches x (a, b)
 210     | x == a = True
 211     | x == b = True
 212     | otherwise = False
 213
 214 anyEdgeTouch xs p = S.foldl' (\t x -> t || (edgeTouches x p)) False xs
 215
 216 edgeOther x (a, b)
 217     | x == a = b
 218     | otherwise = a
 219
 220
 221
 222 contractCave :: ExpandedCaveComplex -> [Position] -> CaveComplex
 223 contractCave expanded startPositions = cavern
 224     where explorers = M.fromList $ zip startPositions $ map intToDigit [0..]
 225           starts = M.union explorers $ _keysE expanded
 226           cavern0 = CaveComplex {_cave = M.empty, _keys = S.fromList $ M.elems $ _keysE expanded}
 227           cavern = M.foldrWithKey (contractFrom expanded) cavern0 starts
 228
 229 contractFrom :: ExpandedCaveComplex -> Position -> Char -> CaveComplex -> CaveComplex
 230 contractFrom expanded startPos startKey cc = cc { _cave = M.union (_cave cc) reachables }
 231     where reachables = reachableFrom [(startPos, edge0)] S.empty expanded' startKey
 232           edge0 = CaveEdge {_keysRequired = S.empty, _distance = 0}
 233           expanded' = expanded {_keysE = M.delete startPos $ _keysE expanded}
 234
 235 reachableFrom :: [(Position, CaveEdge)] -> (S.Set Position) -> ExpandedCaveComplex -> Char -> Cave
 236 reachableFrom [] _closed _expanded _startKey = M.empty
 237 reachableFrom ((here, edge):boundary) closed expanded startKey
 238     | here `S.member` closed = reachableFrom boundary closed expanded startKey
 239     | here `M.member` ks = M.insert edgeKey edge $ reachableFrom boundary closed' expanded startKey
 240     | here `M.member` drs = reachableFrom boundaryD closed' expanded startKey
 241     | otherwise = reachableFrom boundary' closed' expanded startKey
 242     where nbrs0 = S.intersection (_caveE expanded) $ possibleNeighbours here
 243           nbrs = S.difference nbrs0 closed
 244           closed' = S.insert here closed
 245           ks = _keysE expanded
 246           drs = _doors expanded
 247           edgeKey = mkEdgeKey startKey (ks!here)
 248           edge' = edge { _distance = (_distance edge) + 1}
 249           edgeD = edge' {_keysRequired = S.insert (toLower (drs!here)) (_keysRequired edge')}
 250           neighbours = S.map (\n -> (n, edge')) nbrs
 251           neighboursD = S.map (\n -> (n, edgeD)) nbrs
 252           boundary' = boundary ++ (S.toAscList neighbours)
 253           boundaryD = boundary ++ (S.toAscList neighboursD)
 254
 255 possibleNeighbours :: Position -> S.Set Position
 256 possibleNeighbours (r, c) = [(r + 1, c), (r - 1, c), (r, c + 1), (r, c - 1)]
 257
 258
 259 searchCave :: (Explorer e, HasTravelled e Int, HasKeysHeld e Keys) => e -> CaveContext (Maybe (Agendum e))
 260 searchCave explorer =
 261     do agenda <- initAgenda explorer
 262        aStar agenda S.empty
 263
 264 initAgenda :: (Explorer e, HasTravelled e Int, HasKeysHeld e Keys) => e -> CaveContext (Agenda e)
 265 initAgenda explorer =
 266     do cost <- estimateCost explorer
 267        return $ P.singleton cost Agendum { _current = explorer, _trail = Q.empty, _cost = cost}
 268
 269
 270 aStar :: (Explorer e, HasTravelled e Int, HasKeysHeld e Keys) => Agenda e -> ExploredStates e -> CaveContext (Maybe (Agendum e))
 271 aStar agenda closed
 272     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
 273     -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMin agenda) ) False = undefined
 274     | P.null agenda = return Nothing
 275     | otherwise =
 276         do  let (_, currentAgendum) = P.findMin agenda
 277             let reached = _current currentAgendum
 278             nexts <- candidates currentAgendum closed
 279             let newAgenda = foldl' (\q a -> P.insert (_cost a) a q) (P.deleteMin agenda) nexts
 280             reachedGoal <- isGoal reached
 281             if reachedGoal
 282             then return (Just currentAgendum)
 283             else if reached `S.member` closed
 284                  then aStar (P.deleteMin agenda) closed
 285                  else aStar newAgenda (S.insert reached closed)
 286
 287
 288 isGoal :: (Explorer e, HasTravelled e Int, HasKeysHeld e Keys) => e -> CaveContext Bool
 289 isGoal explorer =
 290     do ks <- asks _keys
 291        return $ ks == (explorer ^. keysHeld)
 292
 293
 294 candidates :: (Explorer e, HasTravelled e Int, HasKeysHeld e Keys) => Agendum e -> ExploredStates e -> CaveContext (Q.Seq (Agendum e))
 295 candidates agendum closed =
 296     do  let candidate = _current agendum
 297         let previous = _trail agendum
 298         succs <- successors candidate
 299         let nonloops = Q.filter (\s -> not $ s `S.member` closed) succs
 300         mapM (makeAgendum candidate previous) nonloops
 301
 302 makeAgendum :: (Explorer e, HasTravelled e Int, HasKeysHeld e Keys) => e -> (Q.Seq e) -> e -> CaveContext (Agendum e)
 303 makeAgendum candidate previous new =
 304     do predicted <- estimateCost new
 305        return Agendum { _current = new
 306                       , _trail = candidate <| previous
 307                       , _cost = (new ^. travelled) + predicted
 308                       }