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