adventofcode1611/app/advent11ps.hs

   1 -- Using the idea of canonical representation of buildings from
   2 -- https://andars.github.io/aoc_day11.html by Andrew Foote,
   3 -- plus my extension of represening the pairs as an integer.
   4
   5 -- This version is A* search, using a priority queue for the agenda
   6 -- and HashSets for various collecions.
   7
   8 {-# LANGUAGE DeriveGeneric #-}
   9
  10 module Main(main) where
  11
  12 import GHC.Generics (Generic)
  13
  14 -- import Prelude hiding (length, take, drop)
  15 import Data.List (subsequences, (\\), sort, sortOn, nub, findIndices)
  16 import Data.Ord (comparing)
  17 import Data.Char (isDigit)
  18 import Data.Maybe (fromMaybe)
  19 import qualified Data.PQueue.Prio.Min as P
  20 import Data.Hashable
  21 import qualified Data.HashSet as S
  22 import qualified Data.Sequence as Q
  23 import Data.Sequence ((<|), (|>), (><))
  24 import Data.Foldable (toList, foldr')
  25
  26
  27 data Item = Generator String | Microchip String deriving (Show, Eq, Generic)
  28 instance Hashable Item
  29 type Floor = [Item]
  30 data Building = Building Int [Floor] deriving (Show, Eq, Generic)
  31 instance Hashable Building
  32 data CBuilding = CBuilding Int Integer deriving (Show, Eq, Generic)
  33 instance Hashable CBuilding
  34 type CBuildings = S.HashSet CBuilding
  35 data Agendum = Agendum {current :: Building, trail :: Q.Seq CBuilding, cost :: Int} deriving (Show, Eq)
  36 type Agenda = P.MinPQueue Int Agendum
  37 type Candidates = S.HashSet (Int, Agendum)
  38
  39 instance Ord Item where
  40     compare (Generator a) (Generator b) = compare a b
  41     compare (Microchip a) (Microchip b) = compare a b
  42     compare (Generator _) (Microchip _) = LT
  43     compare (Microchip _) (Generator _) = GT
  44
  45 instance Ord Building where
  46     compare b1 b2 = comparing estimateCost b1 b2
  47
  48 building1 = Building 0 [
  49     (sort [Generator "polonium", Generator "thulium",
  50      Microchip "thulium", Generator "promethium", Generator "ruthenium",
  51      Microchip "ruthenium", Generator "cobalt", Microchip "cobalt"]),
  52     (sort [Microchip "polonium", Microchip "promethium"]),
  53     [],
  54     []
  55     ]
  56
  57 building0 = Building 0 [
  58     (sort [Generator "polonium", Generator "thulium",
  59      Microchip "thulium", Generator "promethium"]),
  60     (sort [Microchip "polonium", Microchip "promethium"]),
  61     [],
  62     []
  63     ]
  64
  65 building2 = Building 0 [
  66     (sort [Generator "polonium", Generator "thulium",
  67      Microchip "thulium", Generator "promethium", Generator "ruthenium",
  68      Microchip "ruthenium", Generator "cobalt", Microchip "cobalt",
  69      Generator "elerium", Microchip "elerium",
  70      Generator "dilithium", Microchip "dilithium"]),
  71     (sort [Microchip "polonium", Microchip "promethium"]),
  72     [],
  73     []
  74     ]
  75
  76
  77 buildingTest = Building 0 [
  78     sort([Microchip "hydrogen", Microchip "lithium"]),
  79     [Generator "hydrogen"],
  80     [Generator "lithium"],
  81     []]
  82
  83 canonical :: Building -> CBuilding
  84 canonical (Building f floors) = CBuilding f (read $ filter (isDigit) $ show $ sort pairs)
  85     where names = nub $ map (\(Generator n) -> n) $ filter (isGenerator) $ concat floors
  86           floorOf (Generator g) = head (findIndices
  87                                                 (\fl -> (Generator g) `elem` fl)
  88                                                 floors)
  89           floorOf (Microchip g) = head (findIndices
  90                                                 (\fl -> (Microchip g) `elem` fl)
  91                                                 floors)
  92           pairs = foldl (\ps n -> (floorOf (Generator n), floorOf (Microchip n)):ps) [] names
  93
  94
  95 main :: IO ()
  96 main = do
  97     part1
  98     part2
  99
 100 part1 :: IO ()
 101 part1 = print $ length $ trail $ fromMaybe (snd $ P.findMin $ initAgenda buildingTest) $ aStar (initAgenda building1) S.empty
 102
 103 part2 :: IO ()
 104 part2 = print $ length $ trail $ fromMaybe (snd $ P.findMin $ initAgenda buildingTest) $ aStar (initAgenda building2) S.empty
 105 initAgenda :: Building -> Agenda
 106 initAgenda b = P.singleton (estimateCost b) Agendum {current = b, trail = Q.empty, cost = estimateCost b}
 107
 108
 109 aStar :: Agenda -> CBuildings -> Maybe Agendum
 110 -- aStar [] _ = Agendum {current=buildingTest, trail=[], cost=0}
 111 aStar agenda closed
 112     | P.null agenda = Nothing
 113     | otherwise =
 114         if isGoal reached then Just currentAgendum
 115         else if creached `S.member` closed
 116             then aStar (P.deleteMin agenda) closed
 117             else aStar newAgenda (S.insert creached closed)
 118         where
 119             (_, currentAgendum) = P.findMin agenda
 120             reached = current currentAgendum
 121             creached = canonical reached
 122             newAgenda = foldr' (\(c, a) q -> P.insert c a q) (P.deleteMin agenda) $ candidates currentAgendum closed
 123             -- newAgenda = P.union (P.deleteMin agenda)
 124             --                     (P.fromList $ toList $ candidates currentAgendum closed)
 125
 126
 127 candidates :: Agendum -> CBuildings -> Q.Seq (Int, Agendum)
 128 candidates agendum closed = newCandidates
 129     where
 130         candidate = current agendum
 131         previous = trail agendum
 132         succs = legalSuccessors $ successors candidate
 133         nonloops = Q.filter (\s -> not $ (canonical s) `S.member` closed) succs
 134         newCandidates = fmap (\a -> (cost a, a)) $ fmap (\n -> makeAgendum n) nonloops
 135         makeAgendum new = Agendum {current = new,
 136                                     trail = (canonical candidate) <| previous,
 137                                     cost = estimateCost new + length previous + 1}
 138
 139 isGoal :: Building -> Bool
 140 isGoal (Building f floors) =
 141     f+1 == height && (all (null) $ take f floors)
 142     where height = length floors
 143
 144 isLegal :: Building -> Bool
 145 isLegal (Building f floors) =
 146     null floor
 147     ||
 148     not (any (isGenerator) floor)
 149     ||
 150     any (safePair) pairs
 151     where floor = floors!!f
 152           pairs = [(i, j) | i <- floor, j <- floor, isGenerator i]
 153           safePair (Generator e, Microchip f) = e == f
 154           safePair (Generator _, Generator _) = False
 155
 156 isGenerator :: Item -> Bool
 157 isGenerator (Generator _) = True
 158 isGenerator (Microchip _) = False
 159
 160 successors :: Building -> (Q.Seq Building)
 161 successors (Building f floors) = Q.fromList [updateBuilding f floors nf is | nf <- nextFloors, is <- items]
 162     where
 163         floor = floors!!f
 164         items = filter (\is -> length is == 1 || length is == 2) $ subsequences floor
 165         nextFloors = if f == 0 then [1]
 166                      else if f+1 == length floors then [f-1]
 167                      else [f+1, f-1]
 168
 169 legalSuccessors :: (Q.Seq Building) -> (Q.Seq Building)
 170 legalSuccessors = Q.filter (isLegal)
 171
 172 updateBuilding :: Int -> [Floor] -> Int -> [Item] -> Building
 173 updateBuilding oldF oldFloors newF items = Building newF newFloors
 174     where newFloors = map (updateFloor) $ zip [0..] oldFloors
 175           updateFloor (f, fl)
 176             | f == oldF = sort $ fl \\ items
 177             | f == newF = sort $ items ++ fl
 178             | otherwise = fl
 179
 180 estimateCost :: Building -> Int
 181 estimateCost (Building _ floors) =
 182     sum $ map (\(c, f) -> c * length f) $ zip [0..] $ reverse floors
 183