{-# 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) -- 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 1 26 27 52] -- [@, a...z, A...Z] data Implicit = Implicit (A.Array Int [Int]) -- edge list (A.UArray (Int, Int) Int) -- distance matrix, with closure taken deriving (Show) implicitGraph :: Map Pos Char -> Pos -> Implicit implicitGraph bd startPos = let posGraph = fst (go startPos Map.empty Set.empty) mapGraph = Map.mapKeys (bd Map.!) (Map.map (Map.mapKeys (bd Map.!)) posGraph) charToNode '@' = 0 charToNode c | isLower c = 1 + ord c - ord 'a' | isUpper c = 1 + 26 + ord c - ord 'A' | otherwise = undefined arrGraph = A.accumArray (const id) [] (0, 2 * 26) [(charToNode from, map charToNode (Map.keys tomap)) | (from, tomap) <- Map.assocs mapGraph] distArr = A.accumArray (const id) (-1) ((0, 0), (2 * 26, 2 * 26)) [((charToNode from, charToNode to), dist) | (from, tomap) <- Map.assocs mapGraph , (to, dist) <- Map.assocs tomap] nodeList = map charToNode (Map.keys mapGraph) in Implicit arrGraph (transitiveClosure nodeList distArr) where 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 Set.foldl' (\(gr, sn) node -> go node gr sn) (graph', seen') newNodes 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 && (c <= 26 || (c - 26) `SIS.member` keys)) (graph A.! at) (nextPearls, nextNonpearls) = partition (\c -> 0 < c && c <= 26 && 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 :: SmallIntSet -> Implicit -> (Int, [Int]) searchBFS allKeys implicit@(Implicit _ distarr) = go 0 (Set.singleton (heuristic 0 SIS.empty, 0, 0, SIS.empty, [])) Map.empty where -- pqueue: f-val, distance, node, keys, key order -- visited: node, keys => distance go :: Int -> Set (Int, Int, Int, SmallIntSet, [Int]) -> Map (Int, SmallIntSet) Int -> (Int, [Int]) go ctr pqueue visited = let ((heurval, dist, curnode, keys, keyorder), newpqueue) = Set.deleteFindMin pqueue reach = reachable implicit keys curnode nextStates = [(dist + stepDist + heuristic stepC stepKeys, dist + stepDist, stepC, stepKeys, stepC : keyorder) | (stepC, stepDist) <- IntMap.assocs reach , let stepKeys = SIS.insert stepC keys -- check that this next state is actually better than we've seen before , maybe True (dist + stepDist <) (Map.lookup (stepC, stepKeys) visited)] visited' = Map.insert (curnode, keys) dist visited pqueue' = newpqueue <> Set.fromList nextStates result = if IntMap.null reach then if heurval == dist then (dist, keyorder) else error ("heurval - dist = " ++ show (heurval - dist) ++ " in terminal state!") else go (ctr + 1) pqueue' visited' in -- (if ctr `rem` 20000 == 0 || IntMap.null reach -- then trace ("go #pqueue=" ++ show (Set.size pqueue) ++ " #visited=" ++ show (Map.size visited) -- ++ " curnode=" ++ show curnode ++ " dist=" ++ show dist ++ " heurval=" ++ show heurval ++ " keys=" ++ show keyorder -- {- ++ " next->" ++ show nextStates -}) -- else id) result heuristic :: Int -> SmallIntSet -> Int heuristic _curnode keys = let remainKeys = allKeys SIS.\\ keys allDists = [distarr A.! (x, y) | x:xs <- tails (SIS.toList remainKeys), y <- xs] distLowerBound = sum (take (SIS.size remainKeys - 1) (sort allDists)) in distLowerBound 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 allKeys = SIS.fromList [ord c - ord 'a' + 1 | c <- Map.elems bd, isLower c] print (fst (searchBFS allKeys imgraph))