1 import Data.List (subsequences, (\\), sort, sortOn, nub, findIndices)
2 import Data.Ord (comparing)
3 import Data.Char (isDigit)
4 import Data.Maybe (fromMaybe)
5 import qualified Data.PQueue.Prio.Min as P
7 data Item = Generator String | Microchip String deriving (Show, Eq)
9 data Building = Building Int [Floor] deriving (Show, Eq)
10 -- data CBuilding = CBuilding Int [(Int, Int)] deriving (Show, Eq)
11 data CBuilding = CBuilding Int Integer deriving (Show, Eq)
12 data Agendum = Agendum {current :: Building, trail :: [CBuilding], cost :: Int}
13 type Agenda = P.MinPQueue Int Agendum
15 instance Ord Item where
16 compare (Generator a) (Generator b) = compare a b
17 compare (Microchip a) (Microchip b) = compare a b
18 compare (Generator _) (Microchip _) = LT
19 compare (Microchip _) (Generator _) = GT
21 instance Ord Building where
22 compare b1 b2 = comparing estimateCost b1 b2
24 building1 = Building 0 [
25 (sort [Generator "polonium", Generator "thulium",
26 Microchip "thulium", Generator "promethium", Generator "ruthenium",
27 Microchip "ruthenium", Generator "cobalt", Microchip "cobalt"]),
28 (sort [Microchip "polonium", Microchip "promethium"]),
33 building0 = Building 0 [
34 (sort [Generator "polonium", Generator "thulium",
35 Microchip "thulium", Generator "promethium"]),
36 (sort [Microchip "polonium", Microchip "promethium"]),
41 building2 = Building 0 [
42 (sort [Generator "polonium", Generator "thulium",
43 Microchip "thulium", Generator "promethium", Generator "ruthenium",
44 Microchip "ruthenium", Generator "cobalt", Microchip "cobalt",
45 Generator "elerium", Microchip "elerium",
46 Generator "dilithium", Microchip "dilithium"]),
47 (sort [Microchip "polonium", Microchip "promethium"]),
53 buildingTest = Building 0 [
54 sort([Microchip "hydrogen", Microchip "lithium"]),
55 [Generator "hydrogen"],
56 [Generator "lithium"],
59 canonical :: Building -> CBuilding
60 canonical (Building f floors) = CBuilding f (read $ filter (isDigit) $ show $ sort pairs)
61 where names = nub $ map (\(Generator n) -> n) $ filter (isGenerator) $ concat floors
62 floorOf (Generator g) = head (findIndices
63 (\fl -> (Generator g) `elem` fl)
65 floorOf (Microchip g) = head (findIndices
66 (\fl -> (Microchip g) `elem` fl)
68 pairs = foldl (\ps n -> (floorOf (Generator n), floorOf (Microchip n)):ps) [] names
77 part1 = print $ length $ trail $ fromMaybe (snd $ P.findMin $ initAgenda buildingTest) $ aStar (initAgenda building1) []
80 part2 = print $ length $ trail $ fromMaybe (snd $ P.findMin $ initAgenda buildingTest) $ aStar (initAgenda building2) []
82 initAgenda :: Building -> Agenda
83 initAgenda b = P.singleton (estimateCost b) Agendum {current = b, trail=[], cost = estimateCost b}
86 aStar :: Agenda -> [CBuilding] -> Maybe Agendum
87 -- aStar [] _ = Agendum {current=buildingTest, trail=[], cost=0}
89 | P.null agenda = Nothing
91 if isGoal reached then Just currentAgendum
92 else if creached `elem` closed
93 then aStar (P.deleteMin agenda) closed
94 else aStar newAgenda (creached:closed)
96 (_, currentAgendum) = P.findMin agenda
97 reached = current currentAgendum
98 creached = canonical reached
99 newAgenda = P.union (P.deleteMin agenda)
100 (P.fromList $ candidates currentAgendum closed)
103 candidates :: Agendum -> [CBuilding] -> [(Int, Agendum)]
104 candidates agendum closed = newCandidates
106 candidate = current agendum
107 previous = trail agendum
108 succs = legalSuccessors $ successors candidate
109 -- nonloops = (succs \\ previous) \\ closed
110 excludable = previous ++ closed
111 nonloops = filter (\s -> not $ (canonical s) `elem` excludable) succs
112 newCandidates = map (\a -> (cost a, a)) $ map (\n -> makeAgendum n) nonloops
113 makeAgendum new = Agendum {current = new,
114 trail = (canonical candidate):previous,
115 cost = estimateCost new + length previous + 1}
117 isGoal :: Building -> Bool
118 isGoal (Building f floors) =
119 f+1 == height && (all (null) $ take f floors)
120 where height = length floors
122 isLegal :: Building -> Bool
123 isLegal (Building f floors) =
126 not (any (isGenerator) floor)
129 where floor = floors!!f
130 pairs = [(i, j) | i <- floor, j <- floor, isGenerator i]
131 safePair (Generator e, Microchip f) = e == f
132 safePair (Generator _, Generator _) = False
134 isGenerator :: Item -> Bool
135 isGenerator (Generator _) = True
136 isGenerator (Microchip _) = False
138 successors :: Building -> [Building]
139 successors (Building f floors) = [updateBuilding f floors nf is | nf <- nextFloors, is <- items]
142 items = filter (\is -> length is == 1 || length is == 2) $ subsequences floor
143 nextFloors = if f == 0 then [1]
144 else if f+1 == length floors then [f-1]
147 legalSuccessors :: [Building] -> [Building]
148 legalSuccessors = filter (isLegal)
150 updateBuilding :: Int -> [Floor] -> Int -> [Item] -> Building
151 updateBuilding oldF oldFloors newF items = Building newF newFloors
152 where newFloors = map (updateFloor) $ zip [0..] oldFloors
154 | f == oldF = sort $ fl \\ items
155 | f == newF = sort $ items ++ fl
158 estimateCost :: Building -> Int
159 estimateCost (Building _ floors) =
160 sum $ map (\(c, f) -> c * length f) $ zip [0..] $ reverse floors