advent11a.hs

   1 import Data.List (subsequences, (\\), sort, sortOn, nub, findIndices)
   2 import Data.Ord (comparing)
   3 import Data.Char (isDigit)
   4
   5 data Item = Generator String | Microchip String deriving (Show, Eq)
   6 type Floor = [Item]
   7 data Building = Building Int [Floor] deriving (Show, Eq)
   8 -- data CBuilding = CBuilding Int [(Int, Int)] deriving (Show, Eq)
   9 data CBuilding = CBuilding Int Integer deriving (Show, Eq)
  10 data Agendum = Agendum {current :: Building, trail :: [CBuilding], cost :: Int}
  11
  12 instance Ord Item where
  13     compare (Generator a) (Generator b) = compare a b
  14     compare (Microchip a) (Microchip b) = compare a b
  15     compare (Generator _) (Microchip _) = LT
  16     compare (Microchip _) (Generator _) = GT
  17
  18 instance Ord Building where
  19     compare b1 b2 = comparing estimateCost b1 b2
  20
  21 building1 = Building 0 [
  22     (sort [Generator "polonium", Generator "thulium",
  23      Microchip "thulium", Generator "promethium", Generator "ruthenium",
  24      Microchip "ruthenium", Generator "cobalt", Microchip "cobalt"]),
  25     (sort [Microchip "polonium", Microchip "promethium"]),
  26     [],
  27     []
  28     ]
  29
  30 building0 = Building 0 [
  31     (sort [Generator "polonium", Generator "thulium",
  32      Microchip "thulium", Generator "promethium"]),
  33     (sort [Microchip "polonium", Microchip "promethium"]),
  34     [],
  35     []
  36     ]
  37
  38 building2 = Building 0 [
  39     (sort [Generator "polonium", Generator "thulium",
  40      Microchip "thulium", Generator "promethium", Generator "ruthenium",
  41      Microchip "ruthenium", Generator "cobalt", Microchip "cobalt",
  42      Generator "elerium", Microchip "elerium",
  43      Generator "dilithium", Microchip "dilithium"]),
  44     (sort [Microchip "polonium", Microchip "promethium"]),
  45     [],
  46     []
  47     ]
  48
  49
  50 buildingTest = Building 0 [
  51     sort([Microchip "hydrogen", Microchip "lithium"]),
  52     [Generator "hydrogen"],
  53     [Generator "lithium"],
  54     []]
  55
  56 canonical :: Building -> CBuilding
  57 -- canonical (Building f floors) = CBuilding f (sort pairs)
  58 canonical (Building f floors) = CBuilding f (read $ filter (isDigit) $ show $ sort pairs)
  59     where names = nub $ map (\(Generator n) -> n) $ filter (isGenerator) $ concat floors
  60           floorOf (Generator g) = head (findIndices
  61                                                 (\fl -> (Generator g) `elem` fl)
  62                                                 floors)
  63           floorOf (Microchip g) = head (findIndices
  64                                                 (\fl -> (Microchip g) `elem` fl)
  65                                                 floors)
  66           pairs = foldl (\ps n -> (floorOf (Generator n), floorOf (Microchip n)):ps) [] names
  67
  68
  69
  70 main :: IO ()
  71 main = do
  72     part1
  73     part2
  74
  75
  76 part1 :: IO ()
  77 part1 = print $ length $ trail $ aStar (initAgenda building1) []
  78 -- part1 = print $ length $ trail $
  79 --             aStar [Agendum {current = building1, trail=[], cost = estimateCost building1}] []
  80
  81 -- part2 :: IO ()
  82 -- part2 = print $ length $ init $ extractJust $ aStar [[building2]] []
  83
  84 part2 :: IO ()
  85 part2 = print $ length $ trail $aStar (initAgenda building2) []
  86
  87 initAgenda :: Building -> [Agendum]
  88 initAgenda b = [Agendum {current = b, trail=[], cost = estimateCost b}]
  89
  90
  91 aStar :: [Agendum] -> [CBuilding] -> Agendum
  92 aStar [] _ = Agendum {current=buildingTest, trail=[], cost=0}
  93 aStar (currentAgendum:agenda) closed =
  94     if isGoal reached then currentAgendum
  95     else if creached `elem` closed
  96         then aStar agenda closed
  97         else aStar newAgenda (creached:closed)
  98     where
  99         reached = current currentAgendum
 100         creached = canonical reached
 101         newAgenda =
 102             -- sortBy (\t1 t2 -> (cost t1) `compare` (cost t2)) $
 103             sortOn (cost) $
 104             agenda ++ (candidates currentAgendum closed)
 105
 106
 107 candidates :: Agendum -> [CBuilding] -> [Agendum]
 108 candidates agendum closed = newCandidates
 109     where
 110         candidate = current agendum
 111         previous = trail agendum
 112         succs = legalSuccessors $ successors candidate
 113         -- nonloops = (succs \\ previous) \\ closed
 114         excludable = previous ++ closed
 115         nonloops = filter (\s -> not $ (canonical s) `elem` excludable) succs
 116         newCandidates = map (\n -> makeAgendum n) nonloops
 117         makeAgendum new = Agendum {current = new,
 118                                     trail = (canonical candidate):previous,
 119                                     cost = estimateCost new + length previous + 1}
 120
 121 isGoal :: Building -> Bool
 122 isGoal (Building f floors) =
 123     f+1 == height && (all (null) $ take f floors)
 124     where height = length floors
 125
 126 isLegal :: Building -> Bool
 127 isLegal (Building f floors) =
 128     null floor
 129     ||
 130     not (any (isGenerator) floor)
 131     ||
 132     any (safePair) pairs
 133     where floor = floors!!f
 134           pairs = [(i, j) | i <- floor, j <- floor, isGenerator i]
 135           safePair (Generator e, Microchip f) = e == f
 136           safePair (Generator _, Generator _) = False
 137
 138 isGenerator :: Item -> Bool
 139 isGenerator (Generator _) = True
 140 isGenerator (Microchip _) = False
 141
 142 successors :: Building -> [Building]
 143 successors (Building f floors) = [updateBuilding f floors nf is | nf <- nextFloors, is <- items]
 144     where
 145         floor = floors!!f
 146         items = filter (\is -> length is == 1 || length is == 2) $ subsequences floor
 147         nextFloors = if f == 0 then [1]
 148                      else if f+1 == length floors then [f-1]
 149                      else [f+1, f-1]
 150
 151 legalSuccessors :: [Building] -> [Building]
 152 legalSuccessors = filter (isLegal)
 153
 154 updateBuilding :: Int -> [Floor] -> Int -> [Item] -> Building
 155 updateBuilding oldF oldFloors newF items = Building newF newFloors
 156     where newFloors = map (updateFloor) $ zip [0..] oldFloors
 157           updateFloor (f, fl)
 158             | f == oldF = sort $ fl \\ items
 159             | f == newF = sort $ items ++ fl
 160             | otherwise = fl
 161
 162 estimateCost :: Building -> Int
 163 estimateCost (Building _ floors) =
 164     sum $ map (\(c, f) -> c * length f) $ zip [0..] $ reverse floors
 165