From 90cdf767b0eaef50cd52b2abb3da5d26e00b9575 Mon Sep 17 00:00:00 2001 From: Daniel Lin Date: Sat, 4 Jan 2025 21:22:09 -0500 Subject: [PATCH] Day 18: Union-Find --- hs/aoc2024.cabal | 1 + hs/src/Day18.hs | 86 ++++++++++++++--------- py/aoc2024/day18.py | 81 +++++++++++++--------- rs/benches/criterion.rs | 8 +-- rs/src/day18.rs | 147 +++++++++++++++++++++++++--------------- rs/src/main.rs | 4 +- 6 files changed, 201 insertions(+), 126 deletions(-) diff --git a/hs/aoc2024.cabal b/hs/aoc2024.cabal index a61b62ed..1458941c 100644 --- a/hs/aoc2024.cabal +++ b/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, diff --git a/hs/src/Day18.hs b/hs/src/Day18.hs index 0257023b..59b703a4 100644 --- a/hs/src/Day18.hs +++ b/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 diff --git a/py/aoc2024/day18.py b/py/aoc2024/day18.py index 52e5581d..9691e54d 100644 --- a/py/aoc2024/day18.py +++ b/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) diff --git a/rs/benches/criterion.rs b/rs/benches/criterion.rs index fb6870fc..e77201ef 100644 --- a/rs/benches/criterion.rs +++ b/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)?; diff --git a/rs/src/day18.rs b/rs/src/day18.rs index 7e01a93b..a0236547 100644 --- a/rs/src/day18.rs +++ b/rs/src/day18.rs @@ -1,66 +1,109 @@ -use std::collections::{BTreeSet, VecDeque}; +use std::collections::VecDeque; +use std::mem::replace; -fn parse(data: &str) -> impl Iterator + use<'_> { +fn parse(data: &str) -> impl Iterator + use<'_> { data.lines().filter_map(|line| { let (x, y) = line.split_once(',')?; - x.parse().ok().zip(y.parse().ok()) + x.parse().ok().zip(y.parse().ok()).map(|pos| (pos, line)) }) } -pub struct Day18 {} - -pub type Default = Day18<70, 1024>; +pub fn part1(data: &str) -> Option { + part1_helper::<70, 1024>(data) +} -impl Day18 { - fn find_path(obstacles: &[(usize, usize)]) -> Option> { - let mut visited = obstacles.iter().copied().collect::>(); - let mut queue: VecDeque<_> = [((0, 0), vec![])].into(); - while let Some((pos @ (x, y), mut path)) = queue.pop_front() { - if !visited.insert(pos) { - continue; - } - path.push(pos); - if x == SIZE && y == SIZE { - return Some(path); - } - if let Some(x) = x.checked_sub(1) { - queue.push_back(((x, y), path.clone())); - } - if let Some(y) = y.checked_sub(1) { - queue.push_back(((x, y), path.clone())); - } - if x < SIZE { - queue.push_back(((x + 1, y), path.clone())); - } - if y < SIZE { - queue.push_back(((x, y + 1), path.clone())); +fn part1_helper(data: &str) -> Option { + let mut visited = vec![false; (SIZE + 1) * (SIZE + 1)]; + visited[0] = true; + for ((x, y), _) in parse(data).take(N) { + visited[x * (SIZE + 1) + y] = true; + } + let mut queue: VecDeque<_> = [((0, 0), 0)].into(); + while let Some(((x, y), t)) = queue.pop_front() { + if x == SIZE && y == SIZE { + return Some(t); + } + for pos @ (x, y) in [ + (x.wrapping_sub(1), y), + (x, y.wrapping_sub(1)), + (x, y + 1), + (x + 1, y), + ] { + if x <= SIZE && y <= SIZE && !visited[x * (SIZE + 1) + y] { + visited[x * (SIZE + 1) + y] = true; + queue.push_back((pos, t + 1)); } } - None } + None +} - pub fn part1(data: &str) -> Option { - Some(Self::find_path(&parse(data).take(N).collect::>())?.len() - 1) - } +pub fn part2(data: &str) -> Option<&str> { + part2_helper::<70>(data) +} - pub fn part2(data: &str) -> Option { - let obstacles = parse(data).collect::>(); - let mut i = 0; - while i < obstacles.len() { - let Some(path) = Self::find_path(&obstacles[..i + 1]) else { - let (x, y) = obstacles[i]; - return Some(format!("{},{}", x, y)); - }; - let path = path.into_iter().collect::>(); - i = obstacles - .iter() - .enumerate() - .skip(i) - .find(|(_, pos)| path.contains(pos))? - .0; +fn part2_helper(data: &str) -> Option<&str> { + let mut obstacles = vec![false; (SIZE + 1) * (SIZE + 1)]; + let candidates = parse(data) + .filter(|((x, y), _)| { + obstacles + .get_mut(x * (SIZE + 1) + y) + .is_some_and(|seen| !replace(seen, true)) + }) + .collect::>(); + let mut sets = (0..(SIZE + 1) * (SIZE + 1)).collect::>(); + fn root(sets: &mut [usize], mut key: usize) -> usize { + let mut value = sets[key]; + while key != value { + let next = sets[value]; + sets[key] = next; + (key, value) = (value, next); + } + value + } + fn union(sets: &mut [usize], i: usize, j: usize) { + let i = root(sets, i); + let j = root(sets, j); + sets[i] = j; + } + for x in 0..=SIZE { + for y in 0..=SIZE { + let i = x * (SIZE + 1) + y; + if obstacles[i] { + continue; + } + for (x, y) in [(x, y + 1), (x + 1, y)] { + if x <= SIZE && y <= SIZE { + let j = x * (SIZE + 1) + y; + if !obstacles[j] { + union(&mut sets, i, j); + } + } + } } - None } + candidates + .into_iter() + .rev() + .find(|((x, y), _)| { + let i = x * (SIZE + 1) + y; + obstacles[i] = false; + for (x, y) in [ + (x.wrapping_sub(1), *y), + (*x, y.wrapping_sub(1)), + (*x, y + 1), + (x + 1, *y), + ] { + if x <= SIZE && y <= SIZE { + let j = x * (SIZE + 1) + y; + if !obstacles[j] { + union(&mut sets, i, j); + } + } + } + root(&mut sets, 0) == root(&mut sets, (SIZE + 1) * (SIZE + 1) - 1) + }) + .map(|(_, line)| line) } #[cfg(test)] @@ -97,15 +140,13 @@ mod tests { 2,0 "}; - type Test = Day18<6, 12>; - #[test] fn part1_examples() { - assert_eq!(Some(22), Test::part1(EXAMPLE)); + assert_eq!(Some(22), part1_helper::<6, 12>(EXAMPLE)); } #[test] fn part2_examples() { - assert_eq!(Some("6,1".to_string()), Test::part2(EXAMPLE)); + assert_eq!(Some("6,1"), part2_helper::<6>(EXAMPLE)); } } diff --git a/rs/src/main.rs b/rs/src/main.rs index 43dc6272..c92d5f02 100644 --- a/rs/src/main.rs +++ b/rs/src/main.rs @@ -159,8 +159,8 @@ fn main() -> anyhow::Result<()> { if args.is_empty() || args.contains("18") { println!("Day 18"); let data = get_day_input(18)?; - println!("{:?}", day18::Default::part1(&data).ok_or(anyhow!("None"))?); - println!("{}", day18::Default::part2(&data).ok_or(anyhow!("None"))?); + println!("{:?}", day18::part1(&data).ok_or(anyhow!("None"))?); + println!("{}", day18::part2(&data).ok_or(anyhow!("None"))?); println!(); }