Updated README to fix a typo and add some clarificaiton about directories
[advent-of-code-16.git] / adventofcode1611 / app / advent11psm.hs
1 -- Using the idea of canonical representation of buildings from
2 -- https://andars.github.io/aoc_day11.html by Andrew Foote,
3 -- plus my extension of represening the pairs as an integer.
4
5 -- This version is A* search, using a priority queue for the agenda,
6 -- Sets for various collecions, and a Map to store the floors in the
7 -- building.
8
9 {-# LANGUAGE DeriveGeneric #-}
10
11 module Main(main) where
12
13 import GHC.Generics (Generic)
14
15 -- import Prelude hiding (length, take, drop)
16 import Data.List (subsequences, (\\), sort, sortOn, nub, findIndices, intercalate)
17 import Data.Ord (comparing)
18 import Data.Char (isDigit)
19 import Data.Maybe (fromMaybe, fromJust)
20 import qualified Data.PQueue.Prio.Min as P
21 import qualified Data.Set as S
22 import qualified Data.Sequence as Q
23 import qualified Data.Map.Strict as M
24 import Data.Hashable
25 import Data.Sequence ((<|), (|>), (><))
26 import Data.Foldable (toList, foldr', foldl', all)
27 import Debug.Trace
28
29 data Item = Generator String | Microchip String deriving (Eq, Generic)
30 type Floor = S.Set Item
31 type Floors = M.Map Int Floor
32 data Building = Building Int Floors deriving (Eq, Ord, Generic)
33 type Buildings = S.Set Building
34 -- data CBuilding = CBuilding Int Integer deriving (Show, Eq, Generic)
35 -- instance Hashable CBuilding
36 -- type CBuildings = S.HashSet CBuilding
37 data Agendum = Agendum {current :: Building, trail :: Q.Seq Building, cost :: Int} deriving (Show, Eq)
38 type Agenda = P.MinPQueue Int Agendum
39 type Candidates = S.Set (Int, Agendum)
40
41 instance Show Item where
42 show (Generator a) = "G" ++ take 2 a
43 show (Microchip a) = "M" ++ take 2 a
44
45 instance Ord Item where
46 compare (Generator a) (Generator b) = compare a b
47 compare (Microchip a) (Microchip b) = compare a b
48 compare (Generator _) (Microchip _) = LT
49 compare (Microchip _) (Generator _) = GT
50
51 -- instance Ord Building where
52 -- compare b1 b2 = comparing estimateCost b1 b2
53
54 instance Show Building where
55 show (Building f floors) = (show f) ++ "<* " ++ (intercalate "; " $ toList $ M.map (showFloor) floors)
56 where showFloor fl = intercalate ", " $ toList $ S.map (show) fl
57
58
59 -- building1 = Building 0 [
60 -- (sort [Generator "polonium", Generator "thulium",
61 -- Microchip "thulium", Generator "promethium", Generator "ruthenium",
62 -- Microchip "ruthenium", Generator "cobalt", Microchip "cobalt"]),
63 -- (sort [Microchip "polonium", Microchip "promethium"]),
64 -- [],
65 -- []
66 -- ]
67
68 building1 = Building 0 (M.fromList
69 [ (0, S.fromList [Generator "polonium", Generator "thulium",
70 Microchip "thulium", Generator "promethium", Generator "ruthenium",
71 Microchip "ruthenium", Generator "cobalt", Microchip "cobalt"])
72 , (1, S.fromList [Microchip "polonium", Microchip "promethium"])
73 , (2, S.empty )
74 , (3, S.empty )
75 ])
76
77
78
79 building0 = Building 0 (M.fromList
80 [ (0, S.fromList [Generator "polonium", Generator "thulium", Microchip "thulium", Generator "promethium"])
81 , (1, S.fromList [Microchip "polonium", Microchip "promethium"])
82 , (2, S.empty )
83 , (3, S.empty )
84 ])
85
86 building2 = Building 0 (M.fromList
87 [ (0, S.fromList [Generator "polonium", Generator "thulium",
88 Microchip "thulium", Generator "promethium", Generator "ruthenium",
89 Microchip "ruthenium", Generator "cobalt", Microchip "cobalt",
90 Generator "dilithium", Microchip "dilithium"])
91 , (1, S.fromList [Microchip "polonium", Microchip "promethium"])
92 , (2, S.empty )
93 , (3, S.empty )
94 ])
95
96 buildingTest = Building 0 (M.fromList
97 [ (0, S.fromList [Microchip "hydrogen", Microchip "lithium"])
98 , (1, S.fromList [Generator "hydrogen"])
99 , (2, S.fromList [Generator "lithium"])
100 , (3, S.empty )
101 ])
102
103
104 main :: IO ()
105 main = do
106 -- part0
107 part1
108 part2
109
110 part0 :: IO ()
111 part0 = print $ length $ trail $ fromMaybe (snd $ P.findMin $ initAgenda buildingTest) $ aStar (initAgenda buildingTest) S.empty
112
113 part1 :: IO ()
114 part1 = print $ length $ trail $ fromMaybe (snd $ P.findMin $ initAgenda buildingTest) $ aStar (initAgenda building1) S.empty
115
116 part2 :: IO ()
117 part2 = print $ length $ trail $ fromMaybe (snd $ P.findMin $ initAgenda buildingTest) $ aStar (initAgenda building2) S.empty
118
119
120 initAgenda :: Building -> Agenda
121 initAgenda b = P.singleton (estimateCost b) Agendum {current = b, trail = Q.empty, cost = estimateCost b}
122
123
124 aStar :: Agenda -> Buildings -> Maybe Agendum
125 -- aStar [] _ = Agendum {current=buildingTest, trail=[], cost=0}
126 aStar agenda closed
127 -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
128 | P.null agenda = Nothing
129 | otherwise =
130 if isGoal reached then Just currentAgendum
131 else if reached `S.member` closed
132 then aStar (P.deleteMin agenda) closed
133 else aStar newAgenda (S.insert reached closed)
134 where
135 (_, currentAgendum) = P.findMin agenda
136 reached = current currentAgendum
137 newAgenda = foldl' (\q a -> P.insert (cost a) a q) (P.deleteMin agenda) $ candidates currentAgendum closed
138
139
140
141 candidates :: Agendum -> Buildings -> Q.Seq Agendum
142 candidates agendum closed = newCandidates
143 where
144 candidate = current agendum
145 previous = trail agendum
146 succs = legalSuccessors $ successors candidate
147 nonloops = Q.filter (\s -> not $ s `S.member` closed) succs
148 newCandidates = fmap (\n -> makeAgendum n) nonloops
149 makeAgendum new = Agendum {current = new,
150 trail = candidate <| previous,
151 cost = estimateCost new + length previous + 1}
152
153 isGoal :: Building -> Bool
154 isGoal (Building f floors) =
155 f+1 == height && (all (S.null) $ M.filterWithKey (\k _ -> k < f) floors)
156 where height = M.size floors
157
158 isLegal :: Building -> Bool
159 isLegal (Building f floors) =
160 null floor
161 ||
162 not (any (isGenerator) floor)
163 ||
164 any (safePair) pairs
165 where floor = fromJust $ M.lookup f floors
166 pairs = [(i, j) | i <- (S.toList floor), j <- (S.toList floor), isGenerator i]
167 safePair (Generator e, Microchip f) = e == f
168 safePair (Generator _, Generator _) = False
169
170 isGenerator :: Item -> Bool
171 isGenerator (Generator _) = True
172 isGenerator (Microchip _) = False
173
174 successors :: Building -> (Q.Seq Building)
175 successors b@(Building f floors) = Q.fromList [updateBuilding b nf is | nf <- nextFloors, is <- items]
176 where
177 floor = fromJust $ M.lookup f floors
178 items = map (S.fromList) $ filter (\is -> length is == 1 || length is == 2) $ subsequences $ toList floor
179 nextFloors = if f == 0 then [1]
180 else if f+1 == length floors then [f-1]
181 else [f+1, f-1]
182
183 legalSuccessors :: (Q.Seq Building) -> (Q.Seq Building)
184 legalSuccessors = Q.filter (isLegal)
185
186 updateBuilding :: Building -> Int -> Floor -> Building
187 updateBuilding (Building oldF oldFloors) newF items = Building newF newFloors
188 where newFloors = M.adjust (\f -> f `S.union` items) newF $ M.adjust (\f -> f `S.difference` items) oldF oldFloors
189
190
191 estimateCost :: Building -> Int
192 estimateCost (Building _ floors) =
193 sum $ map (\(c, f) -> c * S.size f) $ zip [0..] $ reverse $ M.elems floors
194