16-1

1 files changed, 168 insertions, 0 deletions

diff --git a/16-1.hs b/16-1.hs
new file mode 100644
index 0000000..a8640fb
--- /dev/null
+++ b/16-1.hs
@@ -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