18-2

1 files changed, 139 insertions, 0 deletions

diff --git a/18-2.hs b/18-2.hs
new file mode 100644
index 0000000..923d468
--- /dev/null
+++ b/18-2.hs
@@ -0,0 +1,139 @@
+import Control.Monad as M (foldM)
+import Control.Monad.ST
+import Control.Monad.Trans.Class
+import Control.Monad.Trans.Maybe
+import Data.HashMap.Strict as HM
+import Data.HashSet as HS
+import Data.List as L
+import Data.Maybe (fromJust)
+import Data.Tuple as T
+import Prelude
+import Prelude as P (filter, foldl, length, map, read, take)
+import Data.Vector.Mutable as V
+
+type Node = (Int,Int)
+
+type Blocked = HS.HashSet Node
+
+type Visited = HM.HashMap Node Int
+
+type Frontier s = BucketQueue s Node
+
+data BucketQueue s a = BQ Int (V.MVector s [a])
+
+data PathingResult s = Success Int | Fail | Working (Frontier s, Visited)
+
+insert :: BucketQueue s a -> (Int,a) -> ST s (BucketQueue s a)
+insert (BQ top buckets) (prio,x) =
+  V.modify buckets (x:) prio >> return (BQ (min top prio) buckets)
+
+view :: BucketQueue s a -> MaybeT (ST s) (a, BucketQueue s a)
+view (BQ top buckets) = do
+  (x, xs) <- V.read buckets top >>= hoistMaybe . L.uncons
+  lift $ V.write buckets top xs
+  top' <- lift $ findTop top buckets
+  return (x, BQ top' buckets)
+
+findTop :: Int -> V.MVector s [a] -> ST s Int
+findTop n v = do
+  top <- V.read v n
+  case top of
+    [] -> if n < V.length v - 1 then findTop (n+1) v else return n
+    _  -> return n
+
+newBQ :: Int -> ST s (BucketQueue s a)
+newBQ n = BQ n <$> V.replicate n []
+
+neighbours :: Node -> [Node]
+neighbours (x,y) = [ (x+1,y)
+                  , (x-1,y)
+                  , (x,y+1)
+                  , (x,y-1)
+                  ]
+
+expand :: Int -> Blocked -> Node -> [Node]
+expand coordMax blocked = P.filter p . neighbours
+  where
+    p :: Node -> Bool
+    p n@(x,y) = not (HS.member n blocked)
+      && min x y >= 0 && max x y <= coordMax
+
+crowDistance :: Node -> Node -> Int
+crowDistance (x0,y0) (x1,y1) =
+  let a = fromIntegral $ x0 - x1 :: Double
+      b = fromIntegral $ y0 - y1 :: Double
+  in floor $ sqrt (a**2 + b**2)
+
+cost :: Visited -> Node -> Int
+cost = (HM.!)
+
+betterPath :: Visited -> Node -> Node -> Int
+betterPath visited through node
+  | HM.member node visited =
+    min (cost visited node) (cost visited through + 1)
+  | otherwise = cost visited through + 1
+
+heuristic :: Node -> (Node,Int) -> (Node,Int)
+heuristic goal (p,c) = (p, c + crowDistance p goal)
+
+pathingStep :: Int
+            -> Node
+            -> Blocked
+            -> Visited
+            -> Frontier s
+            -> ST s (PathingResult s)
+pathingStep coordMax goal blocked visited frontier =
+  runMaybeT (view frontier) >>= maybe (return Fail) continue
+  where
+    continue (current, rest)
+      | current == goal = return . Success $ cost visited current
+      | otherwise = do
+          let nexts = expand coordMax blocked current
+          let nextsCosts = P.map (betterPath visited current) nexts
+          let visited' = P.foldl (flip $ uncurry HM.insert) visited
+                $ zip nexts nextsCosts
+          frontier' <- M.foldM Main.insert rest
+            . P.map (T.swap . heuristic goal)
+            . P.filter (not . (`HM.member` visited) . fst)
+            $ zip nexts nextsCosts
+          return $ Working (frontier', visited')
+
+findDistance :: Int -> Blocked -> ST s (Maybe Int)
+findDistance coordMax blocked = do
+  let start = (0,0)
+  let goal = (coordMax, coordMax)
+  let visited = HM.insert start 0 HM.empty
+  frontier <- newBQ (coordMax ^ (2 :: Int)) >>= flip Main.insert (0, start)
+
+  loop goal visited frontier
+
+  where
+    loop :: Node -> Visited -> Frontier s -> ST s (Maybe Int)
+    loop g v f = do
+      once <- pathingStep coordMax g blocked v f
+      case once of
+        Success x -> return . Just $ x
+        Fail -> return Nothing
+        Working (f',v') -> loop g v' f'
+
+parsePoint :: String -> Node
+parsePoint s = P.read $ '(' : s ++ ")"
+
+parse :: String -> [Node]
+parse = P.map parsePoint . lines
+
+main :: IO ()
+main = do
+  input <- getContents
+  let blocked = parse input
+  print
+    . fst
+    . fromJust
+    . find ((== Nothing) . snd)
+    . zip [1024..]
+    . P.map (\ x -> runST
+              $ findDistance 70
+              . HS.fromList
+              . P.take x
+              $ blocked)
+    $ [1024..P.length blocked]