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
  17
  18 type Position = (Integer, Integer) -- r, c
  19
  20 type Keys = S.Set Char
  21 type PointOfInterest = M.Map Position Char
  22
  23
  24 data Explorer = Explorer { _position :: Position
  25                          , _keysHeld :: Keys
  26                          } deriving (Eq, Ord, Show)
  27 type ExploredStates = S.Set Explorer
  28
  29 type Cave = S.Set Position
  30 data CaveComplex = CaveComplex { _cave :: Cave
  31                                , _keys :: PointOfInterest
  32                                , _doors :: PointOfInterest
  33                                } deriving (Eq, Ord, Show)
  34 type CaveContext = Reader CaveComplex
  35
  36 data Agendum = Agendum { _current :: Explorer
  37                        , _trail :: Q.Seq Explorer
  38                        , _cost :: Int} deriving (Show, Eq)
  39 type Agenda = P.MinPQueue Int Agendum
  40 type Candidates = S.Set (Int, Agendum)
  41
  42
  43
  44
  45 main :: IO ()
  46 main = do
  47         text <- readFile "data/advent18.txt"
  48         let (cc, explorer) = buildCaveComplex text
  49         -- print cc
  50         -- print explorer
  51         print $ part1 cc explorer
  52
  53 part1 :: CaveComplex -> Explorer -> Int
  54 part1 cave explorer = maybe 0 (( + 1) . _cost ) result
  55     where result = runReader (searchCave explorer) cave
  56
  57 -- -- part1 :: CaveComplex -> Explorer -> Maybe Agendum
  58 -- part1 cave explorer = keySeq (fromJust result)
  59 --     where result = runReader (searchCave explorer) cave
  60
  61
  62 keySeq :: Agendum -> Q.Seq Keys
  63 keySeq agendum = Q.filter (not . S.null) kdiff
  64     where keyss = fmap _keysHeld $ _trail agendum
  65           kdiff = fmap (uncurry S.difference) $ Q.zip ((_keysHeld $ _current agendum) <| keyss) keyss
  66
  67
  68 searchCave :: Explorer -> CaveContext (Maybe Agendum)
  69 searchCave explorer =
  70     do agenda <- initAgenda explorer
  71        aStar agenda S.empty
  72
  73
  74 buildCaveComplex text = foldl' buildCaveRow (cc0, explorer0) $ zip [0..] rows
  75     where cc0 = CaveComplex {_cave = S.empty, _keys = M.empty, _doors = M.empty}
  76           explorer0 = Explorer { _position = (0, 0), _keysHeld = S.empty }
  77           rows = lines text
  78
  79 buildCaveRow (cc, explorer) (r, row) = foldl' (buildCaveCell r) (cc, explorer) $ zip [0..] row
  80
  81 buildCaveCell r (cc, explorer) (c, char)
  82     | char == '.' = (cc', explorer)
  83     | char == '@' = (cc', explorer { _position = here })
  84     | isLower char  = (cc' { _keys = M.insert here char $ _keys cc'}, explorer)
  85     | isUpper char  = (cc' { _doors = M.insert here char $ _doors cc'}, explorer)
  86     | otherwise = (cc, explorer)
  87     where cc' = cc { _cave = S.insert here $ _cave cc }
  88           here = (r, c)
  89
  90
  91
  92
  93 initAgenda :: Explorer -> CaveContext Agenda
  94 initAgenda explorer =
  95     do cost <- estimateCost explorer
  96        return $ P.singleton cost Agendum { _current = explorer, _trail = Q.empty, _cost = cost}
  97
  98
  99 aStar :: Agenda -> ExploredStates -> CaveContext (Maybe Agendum)
 100 -- aStar [] _ = Agendum {current=buildingTest, trail=[], cost=0}
 101 aStar agenda closed
 102     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
 103     -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMin agenda) ) False = undefined
 104     | P.null agenda = return Nothing
 105     | otherwise =
 106         do  let (_, currentAgendum) = P.findMin agenda
 107             let reached = _current currentAgendum
 108             nexts <- candidates currentAgendum closed
 109             let newAgenda = foldl' (\q a -> P.insert (_cost a) a q) (P.deleteMin agenda) nexts
 110             reachedGoal <- isGoal reached
 111             if reachedGoal
 112             then return (Just currentAgendum)
 113             else if reached `S.member` closed
 114                  then aStar (P.deleteMin agenda) closed
 115                  else aStar newAgenda (S.insert reached closed)
 116
 117
 118 isGoal :: Explorer -> CaveContext Bool
 119 isGoal explorer =
 120     do keys <- asks (S.fromList . M.elems . _keys)
 121        return $ keys == _keysHeld explorer
 122
 123
 124 candidates :: Agendum -> ExploredStates -> CaveContext (Q.Seq Agendum)
 125 candidates agendum closed =
 126     do  let candidate = _current agendum
 127         let previous = _trail agendum
 128         succs <- successors candidate
 129         let nonloops = Q.filter (\s -> not $ s `S.member` closed) succs
 130         mapM (makeAgendum candidate previous) nonloops
 131
 132 makeAgendum :: Explorer -> (Q.Seq Explorer) -> Explorer -> CaveContext Agendum
 133 makeAgendum candidate previous new =
 134     do cost <- estimateCost new
 135        return Agendum { _current = new
 136                       , _trail = candidate <| previous
 137                       , _cost = cost + (Q.length previous)
 138                       }
 139
 140 successors :: Explorer -> CaveContext (Q.Seq Explorer)
 141 successors explorer =
 142     do  let here = _position explorer
 143         let locations0 = possibleNeighbours here
 144         cave <- asks _cave
 145         keys <- asks _keys
 146         doors <- asks _doors
 147         let keysHeld = _keysHeld explorer
 148         let locations1 = Q.filter (`S.member` cave) locations0
 149         let locations2 = Q.filter (hasKeyFor doors keysHeld) locations1
 150         return $ fmap (\l -> explorer { _position = l, _keysHeld = pickupKey keys keysHeld l}) locations2
 151
 152
 153 hasKeyFor :: PointOfInterest -> Keys -> Position -> Bool
 154 -- hasKeyFor doors keys here | trace ("hkf: " ++ (intercalate " " [show doors, show keys, show here, show (maybe True (`S.member` keys) $ M.lookup here doors)])) False = undefined
 155 hasKeyFor doors keys here = maybe True keyForDoor $ M.lookup here doors
 156     where keyForDoor d = (toLower d) `S.member` keys
 157     -- if location `M.member` doors
 158     -- then (doors!location) `S.elem` keys
 159     -- else True
 160
 161
 162 pickupKey :: PointOfInterest -> Keys -> Position -> Keys
 163 pickupKey keys held here = maybe held (`S.insert` held) $ M.lookup here keys
 164     -- if here `M.member` keys
 165     -- then S.insert (keys!here) held
 166     -- else held
 167
 168
 169 estimateCost :: Explorer -> CaveContext Int
 170 estimateCost explorer = -- return 0
 171     do keys <- asks _keys
 172        let (r, c) = _position explorer
 173        let unfoundKeys = M.filter (`S.notMember` (_keysHeld explorer)) keys
 174        let minR = minimum $ map fst $ M.keys unfoundKeys
 175        let minC = minimum $ map snd $ M.keys unfoundKeys
 176        let maxR = maximum $ map fst $ M.keys unfoundKeys
 177        let maxC = maximum $ map snd $ M.keys unfoundKeys
 178        let spanR = spanV r minR maxR
 179        let spanC = spanV c minC maxC
 180        if M.null unfoundKeys
 181        then return 0
 182        else return $ fromIntegral (spanR + spanC)
 183        -- return $ sum $ map (manhattan here) $ M.keys unfoundKeys
 184
 185 spanV this minV maxV
 186     | this < minV = maxV - this
 187     | this > maxV = this - minV
 188     -- | this > minV && this < maxV = (this - minV) + (maxV - this)
 189     | otherwise = (this - minV) + (maxV - this)
 190
 191 manhattan :: Position -> Position -> Int
 192 manhattan (r1, c1) (r2, c2) = fromIntegral $ abs (r1 - r2) + abs (c1 - c2)
 193
 194 possibleNeighbours :: Position -> Q.Seq Position
 195 possibleNeighbours (r, c) = [(r + 1, c), (r - 1, c), (r, c + 1), (r, c - 1)]