Day 18: Union-Find

Author: ephemient

hs/aoc2024.cabal

@@ -57,6 +57,7 @@ library
     containers ^>=0.7,
     heap ^>=1.0.4,
     megaparsec ^>=9.7.0,
+    monad-loops ^>=0.4.3,
     parallel ^>=3.2.2.0,
     primitive ^>=0.9.0.0,
     split ^>=0.2.5,

hs/src/Day18.hs

@@ -6,14 +6,19 @@
 module Day18 (part1, part1', part2, part2') where
 
 import Common (readEntire)
-import Data.List.NonEmpty (NonEmpty ((:|)))
-import Data.List.NonEmpty qualified as NonEmpty (cons, toList)
-import Data.Set (Set)
-import Data.Set qualified as Set (empty, fromList, insert, member, notMember)
+import Control.Monad (ap, join, liftM2)
+import Control.Monad.Loops (firstM)
+import Control.Monad.ST (runST)
+import Data.Function (on)
+import Data.Functor (($>))
+import Data.IntSet qualified as IntSet (empty, fromList, insert, notMember)
+import Data.List (scanl')
+import Data.Maybe (listToMaybe)
 import Data.Text (Text)
 import Data.Text qualified as T (lines, stripPrefix)
 import Data.Text.Read (Reader)
 import Data.Text.Read qualified as T (decimal)
+import Data.Vector.Unboxed.Mutable qualified as MV (generate, length, read, write)
 
 coord :: (Integral a) => Reader (a, a)
 coord input = do
@@ -28,39 +33,54 @@ part1 = part1' 70 1024
 part1' :: Int -> Int -> Text -> Either String Int
 part1' size n input = do
   coords <- mapM (readEntire coord) . take n $ T.lines input
-  case go size $ Set.fromList coords of
-    Just path -> Right $ length path - 1
-    Nothing -> Left "no solution"
-
-go :: Int -> Set (Int, Int) -> Maybe (NonEmpty (Int, Int))
-go size visited = go' visited [(0, 0) :| []] []
+  maybe (Left "no solution") Right $ go (IntSet.fromList $ 0 : map index coords) [((0, 0), 0)] []
   where
-    go' visited' (path@(pos@(x, y) :| _) : queue1) queue2
-      | pos `Set.member` visited' = go' visited' queue1 queue2
-      | pos == (size, size) = Just path
-      | otherwise =
-          go' (Set.insert pos visited') queue1 $
-            [ NonEmpty.cons pos' path
-            | pos'@(x', y') <- [(x - 1, y), (x, y - 1), (x, y + 1), (x + 1, y)],
-              0 <= x' && x' <= size && 0 <= y' && y' <= size
-            ]
-              ++ queue2
-    go' _ _ [] = Nothing
-    go' visited' [] queue2 = go' visited' (reverse queue2) []
+    index (x, y) = x * (size + 1) + y
+    go visited (((x, y), t) : queue) queue'
+      | x == size && y == size = Just t
+      | otherwise = go (foldl' (flip $ IntSet.insert . index) visited next) queue $ map (,t + 1) next ++ queue'
+      where
+        next =
+          [ pos'
+          | pos'@(x', y') <- [(x - 1, y), (x, y - 1), (x, y + 1), (x + 1, y)],
+            0 <= x' && x' <= size && 0 <= y' && y' <= size && index pos' `IntSet.notMember` visited
+          ]
+    go _ _ [] = Nothing
+    go visited [] queue = go visited (reverse queue) []
 
 part2 :: Text -> Either String (Int, Int)
 part2 = part2' 70
 
 part2' :: Int -> Text -> Either String (Int, Int)
-part2' size input = mapM (readEntire coord) (T.lines input) >>= go' Set.empty
+part2' size input = do
+  candidates <-
+    reverse
+      . filter (uncurry $ IntSet.notMember . index)
+      . (zip `ap` scanl' (flip $ IntSet.insert . index) IntSet.empty)
+      <$> mapM (readEntire coord) (T.lines input)
+  let obstacles0 = maybe IntSet.empty (uncurry $ IntSet.insert . index) $ listToMaybe candidates
+  maybe (Left "No solution") (Right . fst) $ runST $ do
+    acc <- MV.generate (join (*) $ size + 1) id
+    let root key = MV.read acc key >>= root' key
+        root' key value
+          | key == value = pure value
+          | otherwise = do
+              value' <- root value
+              MV.write acc key value' $> value'
+        union i j = join $ MV.write acc <$> root i <*> root j
+    sequence_
+      [ (union `on` index) pos pos'
+      | pos@(x, y) <- join (liftM2 (,)) [0 .. size],
+        index pos `IntSet.notMember` obstacles0,
+        pos' <- [(x, y + 1) | y < size] ++ [(x + 1, y) | x < size],
+        index pos' `IntSet.notMember` obstacles0
+      ]
+    flip firstM candidates $ \(pos@(x, y), obstacles) -> do
+      sequence_
+        [ (union `on` index) pos pos'
+        | pos' <- [(x - 1, y) | x > 0] ++ [(x, y - 1) | y > 0] ++ [(x, y + 1) | y < size] ++ [(x + 1, y) | x < size],
+          index pos' `IntSet.notMember` obstacles
+        ]
+      (==) <$> root 0 <*> root (MV.length acc - 1)
   where
-    go' visited (candidate : rest) =
-      case go size visited' of
-        Just path ->
-          let path' = Set.fromList $ NonEmpty.toList path
-              (skip, rest') = span (`Set.notMember` path') rest
-           in go' (visited' <> Set.fromList skip) rest'
-        Nothing -> Right candidate
-      where
-        visited' = Set.insert candidate visited
-    go' _ _ = Left "no solution"
+    index (x, y) = x * (size + 1) + y

py/aoc2024/day18.py

@@ -3,7 +3,8 @@
 """
 
 from collections import deque
-from typing import Iterable
+from itertools import islice
+from typing import Generator
 
 SAMPLE_INPUT = """
 5,4
@@ -34,53 +35,69 @@
 """
 
 
-def _parse(data: str) -> list[tuple[int, int]]:
-    return [
-        (int(line[: (i := line.index(","))]), int(line[i + 1 :]))
-        for line in data.splitlines()
-        if "," in line
-    ]
+def _parse(data: str) -> Generator[tuple[int, int]]:
+    for line in data.splitlines():
+        if "," not in line:
+            continue
+        x, y = line.split(",", maxsplit=1)
+        yield int(x), int(y)
 
 
-def findpath(obstacles: Iterable[tuple[int, int]], size: int) -> list[tuple[int, int]]:
-    visited, queue = set(obstacles), deque(([(0, 0)],))
+def part1(data: str, size: int = 70, n: int = 1024) -> int | None:
+    """
+    >>> part1(SAMPLE_INPUT, 6, 12)
+    22
+    """
+    visited, queue = set(islice(_parse(data), n)), deque((((0, 0), 0),))
+    visited.add((0, 0))
     while queue:
-        path = queue.popleft()
-        x, y = pos = path[-1]
+        (x, y), _ = pos, t = queue.popleft()
         if x == size and y == size:
-            return path
-        if pos in visited:
-            continue
-        visited.add(pos)
+            return t
         for pos in ((x - 1, y), (x, y - 1), (x, y + 1), (x + 1, y)):
             x, y = pos
-            if 0 <= x <= size and 0 <= y <= size:
-                queue.append(path + [(x, y)])
+            if 0 <= x <= size and 0 <= y <= size and pos not in visited:
+                visited.add(pos)
+                queue.append((pos, t + 1))
     return None
 
 
-def part1(data: str, size: int = 70, n: int = 1024) -> int:
-    """
-    >>> part1(SAMPLE_INPUT, 6, 12)
-    22
-    """
-    return len(findpath(_parse(data)[:n], size)) - 1
+def _root[T](sets: dict[T, T], key: T) -> T:
+    value = sets.setdefault(key, key)
+    while key != value:
+        sets[key], _ = key, value = value, sets.setdefault(value, value)
+    return value
 
 
-def part2(data: str, size: int = 70) -> str:
+def part2(data: str, size: int = 70) -> str | None:
     """
     >>> part2(SAMPLE_INPUT, 6)
     '6,1'
     """
-    obstacles, i = _parse(data), 0
-    while True:
-        path = findpath(obstacles[: i + 1], size)
-        if path is None:
-            x, y = obstacles[i]
+    obstacles, sets = {}, {}
+    for pos in _parse(data):
+        if pos not in obstacles:
+            obstacles[pos] = None
+    for x in range(size + 1):
+        for y in range(size + 1):
+            pos = x, y
+            if pos in obstacles:
+                continue
+            _root(sets, pos)
+            for pos2 in ((x, y + 1), (x + 1, y)):
+                x2, y2 = pos2
+                if 0 <= x2 <= size and 0 <= y2 <= size and pos2 not in obstacles:
+                    sets[_root(sets, pos)] = _root(sets, pos2)
+    for pos in list(reversed(obstacles.keys())):
+        del obstacles[pos]
+        x, y = pos
+        for pos2 in ((x - 1, y), (x, y - 1), (x, y + 1), (x + 1, y)):
+            x2, y2 = pos2
+            if 0 <= x2 <= size and 0 <= y2 <= size and pos2 not in obstacles:
+                sets[_root(sets, pos)] = _root(sets, pos2)
+        if _root(sets, (0, 0)) == _root(sets, (size, size)):
             return f"{x},{y}"
-        path = set(path)
-        while obstacles[i] not in path:
-            i += 1
+    return None
 
 
 parts = (part1, part2)

rs/benches/criterion.rs

@@ -121,12 +121,8 @@ fn aoc2024_bench(c: &mut Criterion) -> io::Result<()> {
 
     let data = get_day_input(18)?;
     let mut g = c.benchmark_group("day 18");
-    g.bench_function("part 1", |b| {
-        b.iter(|| day18::Default::part1(black_box(&data)))
-    });
-    g.bench_function("part 2", |b| {
-        b.iter(|| day18::Default::part2(black_box(&data)))
-    });
+    g.bench_function("part 1", |b| b.iter(|| day18::part1(black_box(&data))));
+    g.bench_function("part 2", |b| b.iter(|| day18::part2(black_box(&data))));
     g.finish();
 
     let data = get_day_input(19)?;