-- Writeup at https://work.njae.me.uk/2022/12/21/advent-of-code-2022-day-19/
+-- Optimised at https://work.njae.me.uk/2023/07/24/optimising-haskell-example-4/
--- import Debug.Trace
+import Debug.Trace
import AoC
import Data.Text (Text)
import qualified Data.Sequence as Q
import qualified Data.Map.Strict as M
import Data.Map.Strict ((!))
-import Data.MultiSet as MS
+import qualified Data.MultiSet as MS
import Data.Sequence ((|>))
import Data.List
import Data.Maybe
import Control.Parallel.Strategies
import Control.DeepSeq
--- pattern Empty <- (Q.viewl -> Q.EmptyL) where Empty = Q.empty
--- pattern x :< xs <- (Q.viewl -> x Q.:< xs) where (:<) = (Q.<|)
--- pattern xs :> x <- (Q.viewr -> xs Q.:> x) where (:>) = (Q.|>)
data Resource = Ore | Clay | Obsidian | Geode
- deriving (Show, Eq, Ord, Generic)
+ deriving (Show, Eq, Ord, Generic, Enum, Bounded)
instance NFData Resource
type Collection = MS.MultiSet Resource
+hashCollection :: Collection -> Int
+hashCollection c =
+ let k = 200
+ in sum $ zipWith (\r n -> (MS.occur r c) * n) [minBound .. maxBound] $ iterate (* k) 1
+
type Blueprint = M.Map Resource Collection
data TimedBlueprint = TimedBlueprint { getBlueprint :: Blueprint, getTimeLimit :: Int, getMaxRobots :: Collection}
type BlueprintContext = Reader TimedBlueprint
-data SingleSearchState = SingleSearchState
+data SearchState = SearchState
{ _resources :: Collection
, _robots :: Collection
+ , _currentTime :: Int
} deriving (Eq, Show, Ord)
-makeLenses ''SingleSearchState
+makeLenses ''SearchState
+
+type StateHash = (Int, Int, Int)
+hashSearchState :: SearchState -> StateHash
+hashSearchState s = (hashCollection (s ^. resources), hashCollection (s ^. robots), s ^. currentTime)
-instance NFData SingleSearchState where
- rnf (SingleSearchState a b) = rnf a `seq` rnf b `seq` ()
+instance NFData SearchState where
+ rnf (SearchState a b c) = rnf a `seq` rnf b `seq` rnf c `seq` ()
-data Agendum s =
- Agendum { _current :: s
- , _trail :: Q.Seq s
+data Agendum =
+ Agendum { _current :: SearchState
+ , _trail :: Q.Seq SearchState
, _trailBenefit :: Int
, _benefit :: Int
} deriving (Show, Eq, Ord)
makeLenses ''Agendum
-type Agenda s = P.MaxPQueue Int (Agendum s)
-
-type ExploredStates s = S.Set (s, Int, Int)
-
-
-class (Eq s, Ord s, Show s) => SearchState s where
- emptySearchState :: s
- successors :: s -> BlueprintContext (Q.Seq s)
- estimateBenefit :: s -> Int -> BlueprintContext Int
-
-instance SearchState SingleSearchState where
- emptySearchState = SingleSearchState { _resources = MS.empty, _robots = MS.singleton Ore }
-
- successors state =
- do blueprint <- asks getBlueprint
- maxRobots <- asks getMaxRobots
- let buildableRobots = M.keys $ M.filter (\required -> required `MS.isSubsetOf` (state ^. resources)) blueprint
- --- if more bots than needed for making any single bot, don't make more of that bot
- let usefulRobots = MS.foldOccur (\res maxNeeded rs ->
- if (MS.occur res (state ^. robots)) >= maxNeeded
- then Data.List.delete res rs
- else rs
- ) buildableRobots maxRobots
- let madeRobots = [ state & robots %~ MS.insert robot
- & resources %~ ( `MS.difference` (blueprint ! robot) )
- | robot <- usefulRobots
- ]
- let afterBuild = [state] ++ madeRobots
- let afterGather = fmap (\s -> s & resources %~ (MS.union (state ^. robots))) afterBuild
- return $ Q.fromList afterGather
-
-
- estimateBenefit currentState timeElapsed =
- do timeLimit <- asks getTimeLimit
- let timeRemaining = timeLimit - (timeElapsed + 1)
- let currentGeodes = MS.occur Geode (currentState ^. resources)
- let currentRobotsGather = (MS.occur Geode (currentState ^. robots)) * timeRemaining
- let newRobotsGather = (timeRemaining * (timeRemaining + 1)) `div` 2
- return $ currentGeodes + currentRobotsGather + newRobotsGather
+type Agenda = P.MaxPQueue Int Agendum
+type ExploredStates = S.Set StateHash
main :: IO ()
main =
print $ part1 blueprints
print $ part2 blueprints
-part1 :: [(Int, Blueprint)] -> Int
+part1, part2 :: [(Int, Blueprint)] -> Int
part1 blueprints = sum [n * (MS.occur Geode (r ^. resources)) | (n, r) <- results]
- where results = [ (n, _current $ fromJust $ runReader searchSpace (TimedBlueprint blueprint 24 (robotLimits blueprint)) )
- | (n, blueprint) <- blueprints ] :: [(Int, SingleSearchState)]
- robotLimits bp = M.foldl' MS.maxUnion MS.empty bp
--- part1 blueprints = sum [n * (MS.occur Geode (r ^. resources)) | (n, r) <- pResults]
--- where -- results = [ (n, _current $ fromJust $ runReader searchSpace (TimedBlueprint blueprint 24 (robotLimits blueprint)) )
--- -- | (n, blueprint) <- blueprints ] :: [(Int, SingleSearchState)]
--- -- pResults = parMap rdeepseq id results
--- -- pResults = (fmap runABlueprint blueprints) `using` parList rdeepseq
--- pResults = (fmap runABlueprint blueprints) `using` (parList rdeepseq)
--- runABlueprint (n, blueprint) = (n, _current $ fromJust $
--- runReader searchSpace (TimedBlueprint blueprint 24 (robotLimits blueprint)) )
--- robotLimits bp = M.foldl' MS.maxUnion MS.empty bp
-
-part2 :: [(Int, Blueprint)] -> Int
-part2 blueprints = product [MS.occur Geode (r ^. resources) | r <- pResults]
- where results = [ _current $ fromJust $ runReader searchSpace (TimedBlueprint blueprint 32 (robotLimits blueprint))
- | (_, blueprint) <- (take 3 blueprints) ] :: [SingleSearchState]
- pResults = parMap rdeepseq id results
- robotLimits bp = M.foldl' MS.maxUnion MS.empty bp
-
-searchSpace :: SearchState s => BlueprintContext (Maybe (Agendum s))
+-- part1 blueprints = results
+ -- where results = fmap (scoreBlueprint 24) blueprints
+ where results = parMap rdeepseq (scoreBlueprint 24) blueprints
+
+-- part2 :: [(Int, Blueprint)] -> Int
+part2 blueprints = product [MS.occur Geode (r ^. resources) | (_, r) <- results]
+ -- where results = fmap (scoreBlueprint 32) $ take 3 blueprints
+ where results = parMap rdeepseq (scoreBlueprint 32) $ take 3 blueprints
+
+robotLimits :: Blueprint -> Collection
+robotLimits bp = M.foldl' MS.maxUnion MS.empty bp
+
+scoreBlueprint :: Int -> (Int, Blueprint) -> (Int, SearchState)
+scoreBlueprint t (n, bp) = ( n
+ , runReader searchSpace (TimedBlueprint bp t (robotLimits bp))
+ )
+
+searchSpace :: BlueprintContext SearchState
searchSpace =
do agenda <- initAgenda
- -- aStar agenda S.empty
- res <- aStar agenda S.empty
- return (res `seq` res)
+ -- searchAll agenda S.empty emptySearchState
+ result <- aStar agenda S.empty
+ return $ (fromJust result) ^. current
-initAgenda :: SearchState s => BlueprintContext (Agenda s)
+initAgenda :: BlueprintContext Agenda
initAgenda =
do let startState = emptySearchState
- b <- estimateBenefit startState 0
+ b <- estimateBenefit startState
return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
-aStar :: SearchState s => Agenda s -> ExploredStates s -> BlueprintContext (Maybe (Agendum s))
+
+aStar :: Agenda -> ExploredStates -> BlueprintContext (Maybe Agendum)
aStar agenda closed
- -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
- -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " benefit " ++ (show $ fst $ P.findMax agenda) ++ " : elapsed " ++ (show $ Q.length $ _trail $ snd $ P.findMax agenda)) False = undefined
| P.null agenda = return Nothing
| otherwise =
do let (_, currentAgendum) = P.findMax agenda
let reached = currentAgendum ^. current
nexts <- candidates currentAgendum closed
let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
- -- let beamAgenda = P.fromDescList $ P.take 10000 newAgenda -- agenda beam width
- -- let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
- reachedGoal <- isGoal currentAgendum
- let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
- if reachedGoal
+ let cl = hashSearchState reached
+ atTimeLimit <- isTimeLimit currentAgendum
+ if atTimeLimit
then return (Just currentAgendum)
else if (cl `S.member` closed)
then aStar (P.deleteMax agenda) closed
- -- else aStar newAgenda (S.insert cl closed)
else aStar newAgenda (S.insert cl closed)
-candidates :: SearchState s => Agendum s -> ExploredStates s -> BlueprintContext (Q.Seq (Agendum s))
+candidates :: Agendum -> ExploredStates -> BlueprintContext (Q.Seq Agendum)
candidates agendum closed =
do let candidate = agendum ^. current
let previous = agendum ^. trail
- let prevBenefit = agendum ^. trailBenefit
succs <- successors candidate
- succAgs <- mapM (makeAgendum previous prevBenefit) succs
- let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit, Q.length $ s ^. trail) `S.notMember` closed) succAgs
+ succAgs <- mapM (makeAgendum previous) succs
+ let nonloops = Q.filter (\s -> (hashSearchState $ s ^. current) `S.notMember` closed) succAgs
return nonloops
-makeAgendum :: SearchState s => Q.Seq s -> Int -> s -> BlueprintContext (Agendum s)
-makeAgendum previous prevBenefit newState =
- do predicted <- estimateBenefit newState (Q.length previous)
+makeAgendum :: Q.Seq SearchState -> SearchState -> BlueprintContext Agendum
+makeAgendum previous newState =
+ do predicted <- estimateBenefit newState
let newTrail = previous |> newState
- -- let incurred = geodesFound newState
- let incurred = 0
+ let incurred = (MS.occur Geode (newState ^. resources))
return Agendum { _current = newState
, _trail = newTrail
, _trailBenefit = incurred
, _benefit = incurred + predicted
}
-isGoal :: SearchState s => Agendum s -> BlueprintContext Bool
-isGoal agendum =
+isTimeLimit :: Agendum -> BlueprintContext Bool
+isTimeLimit agendum =
do timeLimit <- asks getTimeLimit
- return $ Q.length (agendum ^. trail) == timeLimit
+ return $ (agendum ^. current . currentTime) >= timeLimit
+
+emptySearchState :: SearchState
+emptySearchState = SearchState { _resources = MS.empty, _robots = MS.singleton Ore, _currentTime = 0 }
+
+successors :: SearchState -> BlueprintContext (Q.Seq SearchState)
+successors state =
+ do blueprint <- asks getBlueprint
+ maxRobots <- asks getMaxRobots
+ timeLimit <- asks getTimeLimit
+
+ let robotSuccessors = Q.fromList $ catMaybes $ M.elems $ M.mapWithKey (handleRobot state maxRobots timeLimit) blueprint
+
+ let timeRemaining = timeLimit - (state ^. currentTime)
+ let gathered = MS.foldOccur (\res n acc -> MS.insertMany res (n * timeRemaining) acc)
+ MS.empty
+ (state ^. robots)
+ let delayUntilEnd = (state & currentTime .~ timeLimit
+ & resources %~ (MS.union gathered)
+ )
+ return ( robotSuccessors |> delayUntilEnd )
+
+handleRobot :: SearchState -> Collection -> Int -> Resource -> Collection -> Maybe SearchState
+handleRobot state maxRobots timeLimit robot recipe
+ | sufficientRobots robot state maxRobots = Nothing
+ | otherwise = buildWhenReady robot state recipe timeLimit
+
+-- do I already have enough of this robot?
+sufficientRobots :: Resource -> SearchState -> Collection -> Bool
+sufficientRobots robot state maxRobots =
+ (robot `MS.member` maxRobots)
+ &&
+ ((MS.occur robot (state ^. robots)) >= (MS.occur robot maxRobots))
+
+-- assuming I can't build this robot, how long do I have to wait for the current
+-- robots to gather enough to build it?
+buildDelay :: SearchState -> Collection -> Maybe Int
+buildDelay state recipe
+ -- | MS.null delay = Just 0
+ | all (\r -> MS.member r rbts) (MS.distinctElems shortfall) = Just $ maximum0 $ fmap snd $ MS.toOccurList delay
+ | otherwise = Nothing
+ where shortfall = recipe `MS.difference` (state ^. resources)
+ delay = MS.foldOccur calcOneDelay MS.empty shortfall
+ rbts = state ^. robots
+ calcOneDelay resource count acc =
+ MS.insertMany resource
+ -- (count `div` (MS.occur resource rbts) + 1)
+ (ceiling $ (fromIntegral count) / (fromIntegral $ MS.occur resource rbts))
+ acc
+ maximum0 xs = if (null xs) then 0 else maximum xs
+
+buildWhenReady :: Resource -> SearchState -> Collection -> Int -> Maybe SearchState
+buildWhenReady robot state recipe timeLimit =
+ do waitDelay <- buildDelay state recipe
+ delay <- tooLate (state ^. currentTime) (waitDelay + 1) timeLimit
+ let gathered = MS.foldOccur (\res n acc -> MS.insertMany res (n * delay) acc)
+ MS.empty
+ (state ^. robots)
+ return (state & robots %~ MS.insert robot -- add the robot
+ & resources %~ (MS.union gathered) -- add the gathered resources
+ & resources %~ ( `MS.difference` recipe ) -- remove the resources to build it
+ & currentTime %~ (+ delay)
+ )
+
+tooLate :: Int -> Int -> Int -> Maybe Int
+tooLate current delay timeLimit
+ | (current + delay) <= timeLimit = Just delay
+ | otherwise = Nothing
+
+
+estimateBenefit :: SearchState -> BlueprintContext Int
+estimateBenefit currentState =
+ do timeLimit <- asks getTimeLimit
+ let timeElapsed = currentState ^. currentTime
+ let timeRemaining = timeLimit - timeElapsed
+ let currentRobotsGather = (MS.occur Geode (currentState ^. robots)) * timeRemaining
+ let newRobotsGather = (timeRemaining * (timeRemaining + 1)) `div` 2
+ return $ currentRobotsGather + newRobotsGather
+
-- Parse the input file
successfulParse input =
case parseOnly blueprintsP input of
Left _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
- Right blueprints -> blueprints
\ No newline at end of file
+ Right blueprints -> blueprints
projects / advent-of-code-22.git / blobdiff