Simplify and improve performance

hs/src/Day24.hs

@@ -5,72 +5,43 @@ Description:    <https://adventofcode.com/2022/day/24 Day 24: Blizzard Basin>
 {-# LANGUAGE NondecreasingIndentation, ViewPatterns #-}
 module Day24 (day24a, day24b) where
 
-import Control.Arrow ((&&&), second)
-import Data.Either (partitionEithers)
 import qualified Data.Heap as Heap (FstMinPolicy, insert, singleton, view)
-import Data.List (foldl', transpose)
-import qualified Data.Map as Map ((!), (!?), fromList, fromListWith)
-import Data.Maybe (maybeToList)
-import Data.Semigroup (Arg(Arg), Max(Max), Min(Min))
-import Data.Set (Set)
-import qualified Data.Set as Set (difference, empty, fromDistinctAscList, fromList, insert, lookupMax, lookupMin, member, notMember, toList)
+import Data.List (foldl')
+import qualified Data.Set as Set (insert, notMember, singleton)
 import Data.Text (Text)
-import qualified Data.Text as T (lines, unpack)
+import qualified Data.Text as T (index, length, lines)
 import Data.Tuple (swap)
+import Data.Vector (Vector)
+import qualified Data.Vector as V ((!), fromList, last, length)
 
-parse :: (Enum a, Num a, Ord a) => Text -> Maybe ((a, a), (a, a), [Set (a, a)])
-parse input = (,, spaces) <$> Set.lookupMin startingSpaces <*> Set.lookupMax startingSpaces where
-    deltas = Map.fromList [('^', decY), ('>', incX), ('v', incY), ('<', decX)]
-    (Set.fromDistinctAscList -> startingSpaces, blizzards) = partitionEithers $ do
-        (y, line) <- zip [0..] $ T.lines input
-        (x, char) <- zip [0..] $ T.unpack line
-        if char == '.' then pure $ Left (y, x) else Right <$> do
-        delta <- maybeToList $ deltas Map.!? char
-        pure ((y, x), delta)
-    potentialSpaces = startingSpaces <> Set.fromList (fst <$> blizzards)
-    xy = Map.fromListWith (<>) $ second (Min &&& Max) <$> Set.toList potentialSpaces
-    yx = Map.fromListWith (<>) $ second (Min &&& Max) . swap <$> Set.toList potentialSpaces
-    incX (y, x) | x >= maxX = (y, minX) | otherwise = (y, x + 1)
-      where (Min minX, Max maxX) = xy Map.! y
-    decX (y, x) | x <= minX = (y, maxX) | otherwise = (y, x - 1)
-      where (Min minX, Max maxX) = xy Map.! y
-    incY (y, x) | y >= maxY = (minY, x) | otherwise = (y + 1, x)
-      where (Min minY, Max maxY) = yx Map.! x
-    decY (y, x) | y <= minY = (maxY, x) | otherwise = (y - 1, x)
-      where (Min minY, Max maxY) = yx Map.! x
-    spaces = (potentialSpaces `Set.difference`) . Set.fromList <$> transpose
-        [iterate delta $ delta pos | (pos, delta) <- blizzards]
-
-search :: (Ord a, Ord b) => (a -> [(b, a)]) -> (b, a) -> [a]
-search next = search' Set.empty . Heap.singleton @Heap.FstMinPolicy where
-    search' seen (Heap.view -> Just ((b, a), heap))
-      | Set.member a seen = search' seen heap
-      | otherwise = a : search' seen' heap' where
-            seen' = Set.insert a seen
-            heap' = foldl' (flip Heap.insert) heap .
-                filter (flip Set.notMember seen' . snd) $ next a
-    search' _ _ = []
-
-search' :: (Int, Int) -> (Int, Int) -> [Set (Int, Int)] -> Maybe (Int, [Set (Int, Int)])
-search' start end@(endY, endX) spaces = do
-    Arg n spaces' <- lookup end $ search next (0, (start, Arg 0 spaces))
-    pure (n, spaces')
-  where
-    next ((y, x), Arg n (space:spaces')) =
-      [ (abs (y' - endY) + abs (x' - endX) + n, (pos, Arg (n + 1) spaces'))
-      | pos@(y', x') <- [(y - 1, x), (y, x - 1), (y, x), (y, x + 1), (y + 1, x)]
-      , pos `Set.member` space
-      ]
+search :: Vector Text -> (Int, Int) -> (Int, Int) -> Int -> Maybe Int
+search lines start end@(endX, endY) startTime = go (Set.singleton (start, startTime)) $
+    Heap.singleton @Heap.FstMinPolicy (0, (start, startTime)) where
+    isFree (x, y) time
+      | y < 0 || y >= V.length lines || x < 1 || x > T.length line - 2 = False
+      | y == 0 || y == V.length lines - 1 = T.index line x == '.'
+      | T.index line ((x - 1 + time) `mod` (T.length line - 2) + 1) == '<' = False
+      | T.index line ((x - 1 - time) `mod` (T.length line - 2) + 1) == '>' = False
+      | T.index (lines V.! ((y - 1 + time) `mod` (V.length lines - 2) + 1)) x == '^' = False
+      | T.index (lines V.! ((y - 1 - time) `mod` (V.length lines - 2) + 1)) x == 'v' = False
+      | otherwise = True
+      where line = lines V.! y
+    go seen (Heap.view -> Just ((_, (pos@(x, y), time)), heap))
+      | pos == end = Just time
+      | otherwise = go seen' heap' where
+        choices = filter (\state@(pos, time) -> isFree pos time && Set.notMember state seen) $
+            (, time + 1) <$> [(x - 1, y), (x, y - 1), (x, y), (x, y + 1), (x + 1, y)]
+        seen' = foldl' (flip Set.insert) seen choices
+        heap' = foldl' (flip Heap.insert) heap
+            [(time + abs (x - endX) + abs (y - endY), choice) | choice@((x, y), time) <- choices]
+    go _ _ = Nothing
 
 day24a :: Text -> Maybe Int
-day24a input = do
-    (start, end, spaces) <- parse input
-    fst <$> search' start end spaces
+day24a input = search lines (1, 0) (T.length (V.last lines) - 2, V.length lines - 1) 0 where
+    lines = V.fromList $ T.lines input
 
 day24b :: Text -> Maybe Int
-day24b input = do
-    (start, end, spaces) <- parse input
-    (n1, spaces') <- search' start end spaces
-    (n2, spaces'') <- search' end start spaces'
-    (n3, _) <- search' start end spaces''
-    pure $ n1 + n2 + n3
+day24b input = search lines start end 0 >>= search lines end start >>= search lines start end where
+    lines = V.fromList $ T.lines input
+    start = (1, 0)
+    end = (T.length (V.last lines) - 2, V.length lines - 1)

kt/src/commonMain/kotlin/com/github/ephemient/aoc2022/Day24.kt

@@ -20,29 +20,18 @@ class Day24(private val lines: List<String>) {
 
     private fun search(start: IntPair, end: IntPair, startTime: Int = 0): Int {
         val (endX, endY) = end
+        val seen = mutableSetOf(IndexedValue(startTime, start))
         val queue = PriorityQueue(compareBy(IndexedValue<IndexedValue<IntPair>>::index))
         queue.add(IndexedValue(0, IndexedValue(startTime, start)))
-        val seen = mutableSetOf<IndexedValue<IntPair>>()
         while (!queue.isEmpty()) {
             val entry = queue.remove().value
-            if (!seen.add(entry)) continue
             val (time, position) = entry
             if (position == end) return time
             val (x, y) = position
-            if (isFree(x, y, time + 1)) {
-                queue.add(IndexedValue(time + abs(x - endX) + abs(y - endY), IndexedValue(time + 1, x to y)))
-            }
-            if (isFree(x - 1, y, time + 1)) {
-                queue.add(IndexedValue(time + abs(x - 1 - endX) + abs(y - endY), IndexedValue(time + 1, x - 1 to y)))
-            }
-            if (isFree(x + 1, y, time + 1)) {
-                queue.add(IndexedValue(time + abs(x + 1 - endX) + abs(y - endY), IndexedValue(time + 1, x + 1 to y)))
-            }
-            if (isFree(x, y - 1, time + 1)) {
-                queue.add(IndexedValue(time + abs(x - endX) + abs(y - 1 - endY), IndexedValue(time + 1, x to y - 1)))
-            }
-            if (isFree(x, y + 1, time + 1)) {
-                queue.add(IndexedValue(time + abs(x - endX) + abs(y + 1 - endY), IndexedValue(time + 1, x to y + 1)))
+            for ((x2, y2) in arrayOf(x - 1 to y, x to y - 1, x to y, x to y + 1, x + 1 to y)) {
+                if (!isFree(x2, y2, time + 1)) continue
+                val state = IndexedValue(time + 1, x2 to y2)
+                if (seen.add(state)) queue.add(IndexedValue(time + abs(x2 - endX) + abs(y2 - endY), state))
             }
         }
         throw NoSuchElementException()