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