{-# LANGUAGE TupleSections #-} module Main where import Control.Monad import qualified Data.Array.Unboxed as A import qualified Data.Array.ST as STA import Data.Char import Data.List import qualified Data.IntMap as IntMap import Data.IntMap (IntMap) import qualified Data.Map.Strict as Map import Data.Map.Strict (Map) import Data.Maybe import qualified Data.Set as Set import Data.Set (Set) -- import Debug.Trace import Input import qualified SmallIntSet as SIS import SmallIntSet (SmallIntSet) replaceAtIndex :: Int -> a -> [a] -> [a] replaceAtIndex i _ [] = error ("Index is " ++ show i ++ " items past end in replaceAtIndex") replaceAtIndex 0 val (_:xs) = val : xs replaceAtIndex i val (x:xs) = x : replaceAtIndex (i-1) val xs -- Considers a distance of '-1' to mean 'unconnected'. -- Applies Floyd-Warshall. transitiveClosure :: [Int] -> A.UArray (Int, Int) Int -> A.UArray (Int, Int) Int transitiveClosure nodeList initMatrix = STA.runSTUArray $ do arr <- STA.thaw initMatrix forM_ nodeList $ \k -> forM_ nodeList $ \i -> forM_ nodeList $ \j -> do dij <- STA.readArray arr (i, j) dik <- STA.readArray arr (i, k) dkj <- STA.readArray arr (k, j) if dik /= -1 && dkj /= -1 && (dij == -1 || dik + dkj < dij) then STA.writeArray arr (i, j) (dik + dkj) else return () return arr type Pos = (Int, Int) type Dir = (Int, Int) reachableFrom :: Map Pos Char -> Pos -> Map Pos Int reachableFrom bd startPos = go 0 (Set.singleton startPos) (Set.singleton startPos) Map.empty where go dist seen boundary result = let boundary' = [pos | (x, y) <- Set.toList boundary , (dx, dy) <- [(-1,0), (0,-1), (1,0), (0,1)] , let pos = (x + dx, y + dy) , bd Map.! pos /= '#' , pos `Set.notMember` seen] (things, frees) = partition (\pos -> bd Map.! pos /= '.') boundary' result' = result <> Map.fromList (map (,dist+1) things) in if null frees then result' else go (dist + 1) (seen <> Set.fromList boundary') (Set.fromList frees) result' -- [0 25 26 51 52 ] -- [a...z, A...Z, @...] data Implicit = Implicit Int -- number of starting positions (A.Array Int [Int]) -- edge list (A.UArray (Int, Int) Int) -- distance matrix, with closure taken deriving (Show) codeIsStart, codeIsLower, codeIsUpper :: Int -> Bool codeIsStart n = n >= 2 * 26 codeIsLower n = n < 26 codeIsUpper n = 26 <= n && n < 2 * 26 codeToLower :: Int -> Int codeToLower n = n - 26 -- assumes codeIsUpper implicitGraph :: Map Pos Char -> [Pos] -> Implicit implicitGraph bd startPositions = let nStart = length startPositions nNodes = 2 * 26 + nStart posGraph = fst (goMultiple startPositions Map.empty Set.empty) mapGraph = Map.mapKeys (bd Map.!) (Map.map (Map.mapKeys (bd Map.!)) posGraph) charToNode '@' = charToNode '1' charToNode c | isLower c = ord c - ord 'a' | isUpper c = 26 + ord c - ord 'A' | isDigit c = 26 + 26 + ord c - ord '1' | otherwise = undefined arrGraph = A.accumArray (const id) [] (0, nNodes - 1) [(charToNode from, map charToNode (Map.keys tomap)) | (from, tomap) <- Map.assocs mapGraph] distArr = A.accumArray (const id) (-1) ((0, 0), (nNodes - 1, nNodes - 1)) [((charToNode from, charToNode to), dist) | (from, tomap) <- Map.assocs mapGraph , (to, dist) <- Map.assocs tomap] nodeList = map charToNode (Map.keys mapGraph) in Implicit nStart arrGraph (transitiveClosure nodeList distArr) where goMultiple :: [Pos] -> Map Pos (Map Pos Int) -> Set Pos -> (Map Pos (Map Pos Int), Set Pos) goMultiple curPoses graph seen = foldl' (\(gr, sn) node -> go node gr sn) (graph, seen) curPoses go :: Pos -> Map Pos (Map Pos Int) -> Set Pos -> (Map Pos (Map Pos Int), Set Pos) go curPos graph seen | curPos `Set.member` seen = (graph, seen) | otherwise = let reach = reachableFrom bd curPos newNodes = Map.keysSet reach Set.\\ seen graph' = Map.insert curPos reach graph seen' = Set.insert curPos seen in goMultiple (Set.toList newNodes) graph' seen' reachable :: Implicit -> SmallIntSet -> Int -> IntMap Int reachable (Implicit _ graph distarr) keys start = snd (go 0 (SIS.singleton start) start IntMap.empty) where go dist seen at result = let nexts = filter (\c -> c `SIS.notMember` seen && (codeIsStart c || codeIsLower c || codeToLower c `SIS.member` keys)) (graph A.! at) (nextPearls, nextNonpearls) = partition (\c -> codeIsLower c && c `SIS.notMember` keys) nexts result' = result <> IntMap.fromList [(c, dist + distarr A.! (at, c)) | c <- nextPearls] seen' = seen <> SIS.fromList nexts in -- trace ("reachable-go at=" ++ show at ++ " dist=" ++ show dist ++ " nexts=" ++ show nexts ++ " (allnexts " ++ show (graph A.! at) ++ ")") $ if null nexts then (seen, result') else foldl (\(sn, rs) c -> go (dist + distarr A.! (at, c)) sn c rs) (seen', result') nextNonpearls searchBFS :: Implicit -> (Int, [Int]) searchBFS implicit@(Implicit nstart _ _) = let startNodes = [52 + i | i <- [0..nstart-1]] in go 0 (Set.singleton (0, startNodes, SIS.empty, [])) Map.empty where -- pqueue: f-val, distance, nodes, keys, key order -- visited: nodes, keys => distance go :: Int -> Set (Int, [Int], SmallIntSet, [Int]) -> Map ([Int], SmallIntSet) Int -> (Int, [Int]) go ctr pqueue visited = let ((dist, curnodes, keys, keyorder), newpqueue) = Set.deleteFindMin pqueue reachLists = [reachable implicit keys node | node <- curnodes] nextStates = [(dist + stepDist, stepNodes, stepKeys, stepC : keyorder) | (robotIdx, reach) <- zip [0..] reachLists , (stepC, stepDist) <- IntMap.assocs reach , let stepKeys = SIS.insert stepC keys stepNodes = replaceAtIndex robotIdx stepC curnodes -- check that this next state is actually better than we've seen before , maybe True (dist + stepDist <) (Map.lookup (stepNodes, stepKeys) visited)] visited' = Map.insert (curnodes, keys) dist visited pqueue' = newpqueue <> Set.fromList nextStates result = if maybe True (dist <) (Map.lookup (curnodes, keys) visited) then if all IntMap.null reachLists then (dist, keyorder) else go (ctr + 1) pqueue' visited' else go (ctr + 1) newpqueue visited in -- (if ctr `rem` 20000 == 0 || all IntMap.null reachLists -- then trace ("go ctr=" ++ show ctr ++ " #pqueue=" ++ show (Set.size pqueue) ++ " #visited=" ++ show (Map.size visited) -- ++ " curnodes=" ++ show curnodes ++ " dist=" ++ show dist ++ " keys=" ++ show keyorder -- {- ++ " next->" ++ show nextStates -}) -- else id) result main :: IO () main = do stringbd <- getInput 18 let bd = Map.fromList [((x, y), c) | (y, row) <- zip [0..] stringbd, (x, c) <- zip [0..] row] startpos = fromJust (lookup '@' (map (\(x,y) -> (y,x)) (Map.assocs bd))) let imgraph = implicitGraph bd [startpos] print (fst (searchBFS imgraph)) let (sx, sy) = startpos bd2 = Map.unionWith (const id) bd (Map.fromList [((sx-1,sy-1), '1'), ((sx,sy-1), '#'), ((sx+1,sy-1), '2') ,((sx-1,sy ), '#'), ((sx,sy ), '#'), ((sx+1,sy ), '#') ,((sx-1,sy+1), '3'), ((sx,sy+1), '#'), ((sx+1,sy+1), '4')]) imgraph2 = implicitGraph bd2 [(sx-1,sy-1), (sx+1,sy-1), (sx-1,sy+1), (sx+1,sy+1)] print (fst (searchBFS imgraph2))