app/advent11.hs

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