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