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 $ aStar [[buildingTest]] []
  52 part1 = print $ length $ init $ extractJust $ aStar [[building1]] []
  53
  54 part2 :: IO ()
  55 part2 = print $ length $ init $ extractJust $ aStar [[building2]] []
  56
  57
  58 extractJust :: Maybe [a] -> [a]
  59 extractJust Nothing = []
  60 extractJust (Just x) = x
  61
  62 aStar :: [[Building]] -> [Building] -> Maybe [Building]
  63 aStar [] _ = Nothing
  64 aStar (currentTrail:trails) closed =
  65     if isGoal (head currentTrail) then Just currentTrail
  66     else aStar newAgenda ((head currentTrail): closed)
  67     where newAgenda =
  68             sortBy (\t1 t2 -> (head t1) `compare` (head t2)) $
  69             trails ++ (candidates currentTrail closed)
  70
  71 candidates :: [Building] -> [Building] -> [[Building]]
  72 candidates currentTrail closed = newCandidates
  73     where
  74         (candidate:trail) = currentTrail
  75         succs = legalSuccessors $ successors candidate
  76         nonloops = (succs \\ trail) \\ closed
  77         newCandidates = map (\n -> n:candidate:trail) nonloops
  78
  79 isGoal :: Building -> Bool
  80 isGoal (Building f floors) =
  81     f+1 == height && (all (null) $ take f floors)
  82     where height = length floors
  83
  84 isLegal :: Building -> Bool
  85 isLegal (Building f floors) =
  86     null floor
  87     ||
  88     not (any (isGenerator) floor)
  89     ||
  90     any (safePair) pairs
  91     where floor = floors!!f
  92           pairs = [(i, j) | i <- floor, j <- floor, isGenerator i]
  93           safePair (Generator e, Microchip f) = e == f
  94           safePair (Generator _, Generator _) = False
  95
  96 isGenerator :: Item -> Bool
  97 isGenerator (Generator _) = True
  98 isGenerator (Microchip _) = False
  99
 100 successors :: Building -> [Building]
 101 successors (Building f floors) = [updateBuilding f floors nf is | nf <- nextFloors, is <- items]
 102     where
 103         floor = floors!!f
 104         items = filter (\is -> length is == 1 || length is == 2) $ subsequences floor
 105         nextFloors = if f == 0 then [1]
 106                      else if f+1 == length floors then [f-1]
 107                      else [f+1, f-1]
 108
 109 legalSuccessors :: [Building] -> [Building]
 110 legalSuccessors = filter (isLegal)
 111
 112 updateBuilding :: Int -> [Floor] -> Int -> [Item] -> Building
 113 updateBuilding oldF oldFloors newF items = Building newF newFloors
 114     where newFloors = map (updateFloor) $ zip [0..] oldFloors
 115           updateFloor (f, fl)
 116             | f == oldF = sort $ fl \\ items
 117             | f == newF = sort $ items ++ fl
 118             | otherwise = fl
 119
 120 estimateCost :: Building -> Int
 121 estimateCost (Building _ floors) =
 122     sum $ map (\(c, f) -> c * length f) $ zip [0..] $ reverse floors
 123