diff options
| -rw-r--r-- | 16-1.hs | 168 |
1 files changed, 168 insertions, 0 deletions
@@ -0,0 +1,168 @@ +import Data.Array as A +import Data.Bifunctor (second) +import Data.List as L (find, uncons) +import Data.Maybe (fromJust) +import Data.Tuple (swap) +import Data.Vector as V hiding + ((++), elem, filter, foldl, head, length, map, zip) +import Prelude hiding (replicate) + +data BucketQueue a = BQ Int (V.Vector [a]) deriving Show + +insert :: BucketQueue a -> (Int,a) -> BucketQueue a +insert (BQ top buckets) (prio,x) = + let top' = min top prio + in BQ top' $ V.accum (flip (:)) buckets [ (prio, x) ] + +view :: BucketQueue a -> Maybe (a, BucketQueue a) +view (BQ top buckets) = do + (x, xs) <- buckets V.!? top >>= L.uncons + let buckets' = buckets V.// [ (top, xs) ] + let top' = findTop top buckets' + return (x, BQ top' buckets') + +findTop :: Int -> V.Vector [a] -> Int +findTop n v = case v V.! n of + [] -> if n < length v - 1 then findTop (n+1) v else n + _ -> n + +newBQ :: Int -> BucketQueue a +newBQ n = BQ n $ replicate n [] + +next :: Direction -> (Int,Int) -> (Int,Int) +next North (x,y) = (x,y-1) +next South (x,y) = (x,y+1) +next West (x,y) = (x-1,y) +next East (x,y) = (x+1,y) + +left :: Direction -> Direction +left North = West +left West = South +left South = East +left East = North + +right :: Direction -> Direction +right North = East +right East = South +right South = West +right West = North + +type Node = ((Int,Int), Direction) + +type Frontier = BucketQueue Node + +data Direction = North | South | East | West deriving (Eq,Ix,Ord,Show) + +data Tile = Wall | Tile { untile :: [Direction] } + +type Map = A.Array (Int,Int) Tile + +type Costs = A.Array Node (Maybe Int) + +neighbours :: (Int,Int) -> [(Direction, (Int,Int))] +neighbours (x,y) = [ (North, (x,y-1)) + , (South, (x,y+1)) + , (West, (x-1,y)) + , (East, (x+1,y)) + ] + +parseMap :: [String] -> Map +parseMap s = + let w = length (head s) - 1 + h = length s - 1 + f p = case g p of + '#' -> Wall + _ -> Tile . map fst . filter ((/= '#') . g . snd) $ neighbours p + g (x,y) = s !! y !! x + in array ((0,0),(w,h)) [ ((x,y), f (x,y)) + | x <- [0..w] + , y <- [0..h] + ] + +findStart :: [String] -> (Int,Int) +findStart ls = + let Just p = L.find ((== 'S') . snd) [ ((x,y), ls !! y !! x) + | x <- [0..length (head ls) - 1] + , y <- [0..length ls - 1] + ] + in fst p + +findGoal :: [String] -> (Int,Int) +findGoal ls = + let Just p = L.find ((== 'E') . snd) [ ((x,y), ls !! y !! x) + | x <- [0..length (head ls) - 1] + , y <- [0..length ls - 1] + ] + in fst p + +distance :: (Int,Int) -> (Int,Int) -> Int +distance (x0,y0) (x1,y1) = + let a = fromIntegral $ x0 - x1 + b = fromIntegral $ y0 - y1 + in floor $ sqrt (a^2 + b^2) + +expand :: Map -> (Node, Maybe Int) -> [(Node,Maybe Int)] +expand m ((pos, dir), Just cost) = + [((next dir pos, dir), Just $ cost + 1) | dir `elem` untile (m A.! pos)] + ++ ((\x -> ((pos,x), Just $ cost + 1000)) <$> [left dir, right dir]) + +pathThrough :: Map -> (Node,Maybe Int) -> Node -> (Node,Maybe Int) +pathThrough m through dest = head . filter ((== dest) . fst) $ expand m through + +betterPath :: Map -> (Node,Maybe Int) -> (Node,Maybe Int) -> (Node,Maybe Int) +betterPath m through dest@(node,oldCost) = + let n@(_,newCost) = pathThrough m through node + in if newCost < oldCost + then n + else dest + +cost :: Costs -> Node -> (Node, Maybe Int) +cost costs node = (node, costs A.! node) + +data PathingResult = Success Int | Fail | Working (Frontier,Costs) + +heuristic :: (Int,Int) -> (Node,Int) -> (Node,Int) +heuristic goal (n@(p,_),f) = (n, f + distance goal p) + +pathOnce :: Map -> (Int,Int) -> Frontier -> Costs -> PathingResult +pathOnce m goal frontier costs = maybe Fail continue (view frontier) + where + continue (current@(p,_), rest) + | p == goal = Success $ fromJust (costs A.! current) + | otherwise = + let costs' = costs A.// [ betterPath m (cost costs current) n + | n <- expand m (cost costs current) + ] + frontier' = foldl insert rest + . map (swap . heuristic goal . second fromJust) + . filter ((== Nothing) . (costs A.!) . fst) + $ expand m (cost costs current) + in Working (frontier', costs') + +findPath :: Map -> (Int,Int) -> (Int,Int) -> Int +findPath m start goal = loop frontier costs + where + frontier = insert (newBQ $ 10^6) (0,(start,East)) + + costs = let (lower,upper) = bounds m + in array ((lower,North),(upper,West)) + [ if p == start && d == East + then ((p,d),Just 0) + else ((p,d),Nothing) + | p <- indices m + , d <- [North, South, East, West] + ] + + loop :: Frontier -> Costs -> Int + loop frontier costs = case pathOnce m goal frontier costs of + Success x -> x + Fail -> error "no path" + Working (frontier', costs') -> loop frontier' costs' + +main :: IO () +main = do + input <- lines <$> getContents + let m = parseMap input + let start = findStart input + let goal = findGoal input + print $ findPath m start goal |