advent11p.hs

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