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