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
  26 class (Eq e, Ord e) => ExplorerC e where
  27     successors :: e -> CaveContext (Q.Seq e)
  28     estimateCost :: e -> CaveContext Int
  29     -- positionE :: e -> Position
  30     keysHeldE :: e -> Keys
  31     emptyExplorer :: e
  32
  33
  34 data Explorer1 = Explorer1 { _position1 :: Position
  35                            , _keysHeld1 :: Keys
  36                            } deriving (Eq, Ord, Show)
  37 data Explorer4 = Explorer4 { _position4 :: S.Set Position
  38                            , _keysHeld4 :: Keys
  39                            } deriving (Eq, Ord, Show)
  40 type ExploredStates e = S.Set e
  41
  42 type Cave = S.Set Position
  43 data CaveComplex = CaveComplex { _cave :: Cave
  44                                , _keys :: PointOfInterest
  45                                , _doors :: PointOfInterest
  46                                } deriving (Eq, Ord, Show)
  47 type CaveContext = Reader CaveComplex
  48
  49 data Agendum e = Agendum { _current :: e
  50                        , _trail :: Q.Seq e
  51                        , _cost :: Int} deriving (Show, Eq)
  52 type Agenda e = P.MinPQueue Int (Agendum e)
  53 -- type Candidates e = S.Set (Int, Agendum e)
  54
  55 instance ExplorerC Explorer1 where
  56     successors explorer =
  57         do  let here = _position1 explorer
  58             let locations0 = possibleNeighbours here
  59             cave <- asks _cave
  60             keys <- asks _keys
  61             doors <- asks _doors
  62             let keysHeld = _keysHeld1 explorer
  63             let locations1 = Q.filter (`S.member` cave) locations0
  64             let locations2 = Q.filter (hasKeyFor doors keysHeld) locations1
  65             return $ fmap (\l -> explorer { _position1 = l, _keysHeld1 = pickupKey keys keysHeld l}) locations2
  66
  67     estimateCost explorer = -- return 0
  68         do keys <- asks _keys
  69            let (r, c) = _position1 explorer
  70            let unfoundKeys = M.keysSet $ M.filter (`S.notMember` (_keysHeld1 explorer)) keys
  71            let (minR, maxR, minC, maxC) = bounds $ unfoundKeys
  72            --  = minimum $ map fst $ M.keys unfoundKeys
  73            -- let minC = minimum $ map snd $ M.keys unfoundKeys
  74            -- let maxR = maximum $ map fst $ M.keys unfoundKeys
  75            -- let maxC = maximum $ map snd $ M.keys unfoundKeys
  76            let spanR = spanV r minR maxR
  77            let spanC = spanV c minC maxC
  78            if S.null unfoundKeys
  79            then return 0
  80            else return $ fromIntegral (spanR + spanC)
  81            -- return $ sum $ map (manhattan here) $ M.keys unfoundKeys
  82
  83     -- positionE = _position1
  84     keysHeldE = _keysHeld1
  85     emptyExplorer = Explorer1 { _position1 = (0, 0), _keysHeld1 = S.empty }
  86
  87 instance ExplorerC Explorer4 where
  88     successors explorer =
  89         do  let rawHeres = _position4 explorer
  90             let heres = setToSeq $ allSplits rawHeres
  91             let locations0 = over (traversed . _1) possibleNeighbours heres
  92             cave <- asks _cave
  93             keys <- asks _keys
  94             doors <- asks _doors
  95             let keysHeld = _keysHeld4 explorer
  96             let locations1 = over (traversed . _1) (Q.filter (`S.member` cave)) locations0
  97             let locations2 = over (traversed . _1) (Q.filter (hasKeyFor doors keysHeld)) locations1
  98             let locations3 = fmap (\(ls, hs) -> fmap (\l -> (l, hs)) ls) locations2
  99             let locations4 = foldl1 (><) locations3
 100             return $ fmap (\(l, hs) -> explorer { _position4 = S.insert l hs, _keysHeld4 = pickupKey keys keysHeld l}) locations4
 101
 102     estimateCost explorer = -- return 0
 103         do keys <- asks _keys
 104            let unfoundKeys = M.keysSet $ M.filter (`S.notMember` (_keysHeld4 explorer)) keys
 105            let (minR, maxR, minC, maxC) = bounds unfoundKeys
 106            let (minDR, maxDR, minDC, maxDC) = bounds $ _position4 explorer
 107            let dr = abs (minR - minDR) + abs (maxR - maxDR)
 108            let dc = abs (minC - minDC) + abs (maxC - maxDC)
 109            if S.null unfoundKeys
 110            then return 0
 111            else return $ fromIntegral (dr + dc)
 112            -- return $ sum $ map (manhattan here) $ M.keys unfoundKeys
 113
 114     -- positionE = _position1
 115     keysHeldE = _keysHeld4
 116     emptyExplorer = Explorer4 { _position4 = S.fromList $ replicate 4 (0, 0), _keysHeld4 = S.empty }
 117
 118
 119
 120 main :: IO ()
 121 main = do
 122         text <- readFile "data/advent18.txt"
 123         let (cc, explorer) = buildCaveComplex text
 124         -- print cc
 125         -- print explorer
 126         print $ part1 cc explorer
 127         print $ part2 cc explorer
 128
 129 part1 :: ExplorerC e => CaveComplex -> e -> Int
 130 part1 cave explorer = maybe 0 (( + 1) . _cost ) result
 131     where result = runReader (searchCave explorer) cave
 132
 133 -- -- part1 :: CaveComplex -> Explorer -> Maybe Agendum
 134 -- part1 cave explorer = keySeq (fromJust result)
 135 --     where result = runReader (searchCave explorer) cave
 136
 137
 138 part2 ::  CaveComplex -> Explorer1 -> Int
 139 part2 caveComplex0 explorer1 = maybe 0 (( + 1) . _cost ) result
 140     where
 141         (re, ce) = _position1 explorer1
 142         cave0 = _cave caveComplex0
 143         cave = cave0 `S.difference` [(re, ce), (re + 1, ce), (re - 1, ce), (re, ce + 1), (re, ce - 1)]
 144         caveComplex = caveComplex0 {_cave = cave}
 145         explorer = Explorer4 {_position4 = [(re + 1, ce + 1), (re - 1, ce + 1), (re + 1, ce - 1), (re - 1, ce - 1)], _keysHeld4 = S.empty }
 146         result = runReader (searchCave explorer) caveComplex
 147
 148 keySeq :: ExplorerC e => (Agendum e) -> Q.Seq Keys
 149 keySeq agendum = Q.filter (not . S.null) kdiff
 150     where keyss = fmap keysHeldE $ _trail agendum
 151           kdiff = fmap (uncurry S.difference) $ Q.zip ((keysHeldE $ _current agendum) <| keyss) keyss
 152
 153
 154 searchCave :: ExplorerC e => e -> CaveContext (Maybe (Agendum e))
 155 searchCave explorer =
 156     do agenda <- initAgenda explorer
 157        aStar agenda S.empty
 158
 159
 160 buildCaveComplex text = foldl' buildCaveRow (cc0, explorer0) $ zip [0..] rows
 161     where cc0 = CaveComplex {_cave = S.empty, _keys = M.empty, _doors = M.empty}
 162           explorer0 = emptyExplorer -- Explorer { _position = (0, 0), _keysHeld = S.empty }
 163           rows = lines text
 164
 165 buildCaveRow (cc, explorer) (r, row) = foldl' (buildCaveCell r) (cc, explorer) $ zip [0..] row
 166
 167 buildCaveCell r (cc, explorer) (c, char)
 168     | char == '.' = (cc', explorer)
 169     | char == '@' = (cc', explorer { _position1 = here })
 170     | isLower char  = (cc' { _keys = M.insert here char $ _keys cc'}, explorer)
 171     | isUpper char  = (cc' { _doors = M.insert here char $ _doors cc'}, explorer)
 172     | otherwise = (cc, explorer)
 173     where cc' = cc { _cave = S.insert here $ _cave cc }
 174           here = (r, c)
 175
 176
 177 initAgenda :: ExplorerC e => e -> CaveContext (Agenda e)
 178 initAgenda explorer =
 179     do cost <- estimateCost explorer
 180        return $ P.singleton cost Agendum { _current = explorer, _trail = Q.empty, _cost = cost}
 181
 182
 183 aStar :: ExplorerC e => Agenda e -> ExploredStates e -> CaveContext (Maybe (Agendum e))
 184 -- aStar [] _ = Agendum {current=buildingTest, trail=[], cost=0}
 185 aStar agenda closed
 186     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
 187     -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMin agenda) ) False = undefined
 188     | P.null agenda = return Nothing
 189     | otherwise =
 190         do  let (_, currentAgendum) = P.findMin agenda
 191             let reached = _current currentAgendum
 192             nexts <- candidates currentAgendum closed
 193             let newAgenda = foldl' (\q a -> P.insert (_cost a) a q) (P.deleteMin agenda) nexts
 194             reachedGoal <- isGoal reached
 195             if reachedGoal
 196             then return (Just currentAgendum)
 197             else if reached `S.member` closed
 198                  then aStar (P.deleteMin agenda) closed
 199                  else aStar newAgenda (S.insert reached closed)
 200
 201
 202 isGoal :: ExplorerC e => e -> CaveContext Bool
 203 isGoal explorer =
 204     do keys <- asks (S.fromList . M.elems . _keys)
 205        return $ keys == keysHeldE explorer
 206
 207
 208 candidates :: ExplorerC e => Agendum e -> ExploredStates e -> CaveContext (Q.Seq (Agendum e))
 209 candidates agendum closed =
 210     do  let candidate = _current agendum
 211         let previous = _trail agendum
 212         succs <- successors candidate
 213         let nonloops = Q.filter (\s -> not $ s `S.member` closed) succs
 214         mapM (makeAgendum candidate previous) nonloops
 215
 216 makeAgendum :: ExplorerC e => e -> (Q.Seq e) -> e -> CaveContext (Agendum e)
 217 makeAgendum candidate previous new =
 218     do cost <- estimateCost new
 219        return Agendum { _current = new
 220                       , _trail = candidate <| previous
 221                       , _cost = cost + (Q.length previous)
 222                       }
 223
 224
 225
 226 hasKeyFor :: PointOfInterest -> Keys -> Position -> Bool
 227 -- 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
 228 hasKeyFor doors keys here = maybe True keyForDoor $ M.lookup here doors
 229     where keyForDoor d = (toLower d) `S.member` keys
 230     -- if location `M.member` doors
 231     -- then (doors!location) `S.elem` keys
 232     -- else True
 233
 234
 235 pickupKey :: PointOfInterest -> Keys -> Position -> Keys
 236 pickupKey keys held here = maybe held (`S.insert` held) $ M.lookup here keys
 237     -- if here `M.member` keys
 238     -- then S.insert (keys!here) held
 239     -- else held
 240
 241
 242 spanV this minV maxV
 243     | this < minV = maxV - this
 244     | this > maxV = this - minV
 245     -- | this > minV && this < maxV = (this - minV) + (maxV - this)
 246     | otherwise = (this - minV) + (maxV - this)
 247
 248 manhattan :: Position -> Position -> Int
 249 manhattan (r1, c1) (r2, c2) = fromIntegral $ abs (r1 - r2) + abs (c1 - c2)
 250
 251 possibleNeighbours :: Position -> Q.Seq Position
 252 possibleNeighbours (r, c) = [(r + 1, c), (r - 1, c), (r, c + 1), (r, c - 1)]
 253
 254 bounds :: S.Set Position -> (Integer, Integer, Integer, Integer)
 255 bounds points = (minR, maxR, minC, maxC)
 256     where  minR = S.findMin $ S.map fst points
 257            minC = S.findMin $ S.map snd points
 258            maxR = S.findMax $ S.map fst points
 259            maxC = S.findMax $ S.map snd points
 260
 261
 262 allSplits :: Ord a => S.Set a -> S.Set (a, S.Set a)
 263 allSplits xs = S.map (\x -> (x, S.delete x xs)) xs
 264
 265 setToSeq :: Ord a => S.Set a -> Q.Seq a
 266 setToSeq = S.foldl (|>) Q.empty