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