{-# LANGUAGE MultiWayIf, TupleSections, ScopedTypeVariables #-} module Main where import qualified Data.Array as A import Data.Char import Data.List import qualified Data.Map.Strict as Map import Data.Maybe import qualified Data.Set as Set import Input shortestPath :: forall node. Ord node => (node -> [node]) -> (node -> Int) -> node -> node -> Int shortestPath nextf heuristic startnode target = go 0 (Set.singleton (heuristic startnode, 0, startnode)) Map.empty where go :: Int -> Set.Set (Int, Int, node) -> Map.Map node Int -> Int go ctr pqueue visited = let ((_, dist, node), newpqueue) = Set.deleteFindMin pqueue isBetter n = maybe True (> dist + 1) (Map.lookup n visited) nexts = filter isBetter (nextf node) pqueue' = newpqueue <> Set.fromList [(dist + 1 + heuristic n, dist + 1, n) | n <- nexts] visited' = Map.insert node dist visited in if node == target then dist else go (ctr + 1) pqueue' visited' main :: IO () main = do input <- getInput 20 let h = length input w = length (head input) bd = A.array ((0, 0), (w - 1, h - 1)) [((x, y), c) | (y, row) <- zip [0..] input, (x, c) <- zip [0..] row] bdget x y | 0 <= x, x < w, 0 <= y, y < h = bd A.! (x, y) | otherwise = ' ' labels = [(name, loc) | y <- [0..h-1] , x <- [0..w-1] , let c1 = bd A.! (x, y) , isUpper c1 , (c2, locpts) <- filter (isUpper . fst) [(bdget x (y+1), [(x, y-1), (x, y+2)]) ,(bdget (x+1) y, [(x-1, y), (x+2, y)])] , let name = [c1, c2] , loc <- filter ((== '.') . uncurry bdget) locpts] label2pos = Map.fromListWith (++) (map (fmap pure) labels) pos2label = Map.fromList (map (\(x, y) -> (y, x)) labels) pos2other = Map.fromList [(pos, other) | (pos, lab) <- Map.assocs pos2label , let poses = label2pos Map.! lab , other <- poses \\ [pos]] isOuter (x, y) = x == 2 || y == 2 || x == w - 3 || y == h - 3 donutWidth = minimum (map fst (filter (not . isOuter) (Map.keys pos2label))) - 3 let mazeNexts (x, y) = maybeToList (Map.lookup (x, y) pos2other) ++ [(x', y') | (dx, dy) <- [(-1,0), (0,-1), (1,0), (0,1)] , let (x', y') = (x + dx, y + dy) , bdget x' y' == '.'] print $ shortestPath mazeNexts (const 0) (head $ label2pos Map.! "AA") (head $ label2pos Map.! "ZZ") let mazeNexts2 (lvl, (x, y)) = [(lvl', other) | other <- maybeToList (Map.lookup (x, y) pos2other) , let lvl' = lvl + (if isOuter other then 1 else -1) , lvl' >= 0] ++ [(lvl, (x', y')) | (dx, dy) <- [(-1,0), (0,-1), (1,0), (0,1)] , let (x', y') = (x + dx, y + dy) , bdget x' y' == '.'] print $ shortestPath mazeNexts2 (\(lvl, _) -> donutWidth * lvl) (0, head $ label2pos Map.! "AA") (0, head $ label2pos Map.! "ZZ")