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