eb6219ff6ad85afd8179b54e67b1aca4b4b9b8e5
[advent-of-code-19.git] / 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