Removed language extension
[advent-of-code-16.git] / advent11.hs
1 import Data.List (subsequences, (\\), sort, sortBy)
2 import Data.Ord (comparing)
3
4 data Item = Generator String | Microchip String deriving (Show, Eq)
5 type Floor = [Item]
6 data Building = Building Int [Floor] deriving (Show, Eq)
7
8 instance Ord Item where
9 compare (Generator a) (Generator b) = compare a b
10 compare (Microchip a) (Microchip b) = compare a b
11 compare (Generator _) (Microchip _) = LT
12 compare (Microchip _) (Generator _) = GT
13
14 instance Ord Building where
15 compare b1 b2 = comparing estimateCost b1 b2
16
17 building1 = Building 0 [
18 (sort [Generator "polonium", Generator "thulium",
19 Microchip "thulium", Generator "promethium", Generator "ruthenium",
20 Microchip "ruthenium", Generator "cobalt", Microchip "cobalt"]),
21 (sort [Microchip "polonium", Microchip "promethium"]),
22 [],
23 []
24 ]
25
26 building2 = Building 0 [
27 (sort [Generator "polonium", Generator "thulium",
28 Microchip "thulium", Generator "promethium", Generator "ruthenium",
29 Microchip "ruthenium", Generator "cobalt", Microchip "cobalt",
30 Generator "elerium", Microchip "elerium",
31 Generator "dilithium", Microchip "dilithium"]),
32 (sort [Microchip "polonium", Microchip "promethium"]),
33 [],
34 []
35 ]
36
37
38 buildingTest = Building 0 [
39 sort([Microchip "hydrogen", Microchip "lithium"]),
40 [Generator "hydrogen"],
41 [Generator "lithium"],
42 []]
43
44 main :: IO ()
45 main = do
46 part1
47 part2
48
49
50 part1 :: IO ()
51 -- part1 = print $ length $ init $ extractJust $ hillClimb [[buildingTest]] []
52 part1 = print $ length $ init $ extractJust $ hillClimb [[building1]] []
53 -- part1 = print $ length $ init $ extractJust $ aStar [[building1]] []
54
55 part2 :: IO ()
56 part2 = print $ length $ init $ extractJust $ hillClimb [[building2]] []
57
58
59 extractJust :: Maybe [a] -> [a]
60 extractJust Nothing = []
61 extractJust (Just x) = x
62
63 hillClimb :: [[Building]] -> [Building] -> Maybe [Building]
64 hillClimb [] _ = Nothing
65 hillClimb (currentTrail:trails) closed =
66 if isGoal (head currentTrail) then Just currentTrail
67 else hillClimb newAgenda ((head currentTrail): closed)
68 where newAgenda =
69 sortBy (\t1 t2 -> (head t1) `compare` (head t2)) $
70 trails ++ (candidates currentTrail closed)
71
72 aStar :: [[Building]] -> [Building] -> Maybe [Building]
73 aStar [] _ = Nothing
74 aStar (currentTrail:trails) closed =
75 if isGoal (head currentTrail) then Just currentTrail
76 else aStar newAgenda ((head currentTrail): closed)
77 where newAgenda =
78 sortBy (\t1 t2 -> (trailCost t1) `compare` (trailCost t2)) $
79 trails ++ (candidates currentTrail closed)
80 trailCost t = estimateCost (head t) + length t - 1
81
82
83 candidates :: [Building] -> [Building] -> [[Building]]
84 candidates currentTrail closed = newCandidates
85 where
86 (candidate:trail) = currentTrail
87 succs = legalSuccessors $ successors candidate
88 nonloops = (succs \\ trail) \\ closed
89 newCandidates = map (\n -> n:candidate:trail) nonloops
90
91 isGoal :: Building -> Bool
92 isGoal (Building f floors) =
93 f+1 == height && (all (null) $ take f floors)
94 where height = length floors
95
96 isLegal :: Building -> Bool
97 isLegal (Building f floors) =
98 null floor
99 ||
100 not (any (isGenerator) floor)
101 ||
102 any (safePair) pairs
103 where floor = floors!!f
104 pairs = [(i, j) | i <- floor, j <- floor, isGenerator i]
105 safePair (Generator e, Microchip f) = e == f
106 safePair (Generator _, Generator _) = False
107
108 isGenerator :: Item -> Bool
109 isGenerator (Generator _) = True
110 isGenerator (Microchip _) = False
111
112 successors :: Building -> [Building]
113 successors (Building f floors) = [updateBuilding f floors nf is | nf <- nextFloors, is <- items]
114 where
115 floor = floors!!f
116 items = filter (\is -> length is == 1 || length is == 2) $ subsequences floor
117 nextFloors = if f == 0 then [1]
118 else if f+1 == length floors then [f-1]
119 else [f+1, f-1]
120
121 legalSuccessors :: [Building] -> [Building]
122 legalSuccessors = filter (isLegal)
123
124 updateBuilding :: Int -> [Floor] -> Int -> [Item] -> Building
125 updateBuilding oldF oldFloors newF items = Building newF newFloors
126 where newFloors = map (updateFloor) $ zip [0..] oldFloors
127 updateFloor (f, fl)
128 | f == oldF = sort $ fl \\ items
129 | f == newF = sort $ items ++ fl
130 | otherwise = fl
131
132 estimateCost :: Building -> Int
133 estimateCost (Building _ floors) =
134 sum $ map (\(c, f) -> c * length f) $ zip [0..] $ reverse floors
135