2024 Day 16 Part 2

2024/src/aoc.zig

@@ -129,7 +129,8 @@ pub fn shortestPath(
     allocator: std.mem.Allocator,
     ctx: State.Context,
     start: State,
-) !?u64 {
+    comptime withGraph: bool,
+) !?(if (withGraph) struct { u64, std.AutoHashMap(State, AutoHashSet(State)) } else u64) {
     const Entry = struct {
         state: State,
         cost: u64,
@@ -149,8 +150,12 @@ pub fn shortestPath(
         .cost = 0,
     });
 
-    var prev = std.AutoHashMap(State, State).init(allocator);
-    defer prev.deinit();
+    var prev = std.AutoHashMap(State, AutoHashSet(State)).init(allocator);
+    defer if (!withGraph) {
+        prev.deinit();
+        var it = prev.valueIterator();
+        while (it.next()) |set| set.deinit();
+    };
 
     var dist = std.AutoHashMap(State, u64).init(allocator);
     defer dist.deinit();
@@ -166,8 +171,12 @@ pub fn shortestPath(
         }
 
         if (u.isGoal(ctx)) {
+            if (withGraph) {
+                return .{ cost, prev };
+            } else {
                 return cost;
             }
+        }
 
         var neighbors: std.ArrayList(State.Neighbor) = try u.neighbors(ctx);
         defer neighbors.deinit();
@@ -177,12 +186,17 @@ pub fn shortestPath(
             const alt = cost + extra;
 
             if (dist.get(v)) |dv| {
-                if (alt >= dv) {
+                if (alt > dv) {
                     continue;
                 }
             }
 
-            try prev.put(v, u);
+            var e = try prev.getOrPut(v);
+            if (!e.found_existing) {
+                e.value_ptr.* = AutoHashSet(State).init(allocator);
+            }
+            try e.value_ptr.put(u);
+
             try dist.put(v, alt);
             try queue.add(Entry{
                 .state = v,

2024/src/day16.zig

@@ -23,11 +23,11 @@ pub fn main() !void {
     var bw = std.io.bufferedWriter(std.io.getStdOut().writer());
     const stdout = bw.writer();
 
-    try stdout.print("{d}\n", .{try part1(allocator, grid)});
+    try stdout.print("{d} {d}\n", try search(allocator, grid));
     try bw.flush();
 }
 
-fn part1(allocator: Allocator, grid: Grid) !u64 {
+fn search(allocator: Allocator, grid: Grid) !struct { u64, u64 } {
     const pos, const goal = v: {
         var start: Pos = undefined;
         var end: Pos = undefined;
@@ -116,7 +116,39 @@ fn part1(allocator: Allocator, grid: Grid) !u64 {
         .goal = goal,
     };
 
-    return (try aoc.shortestPath(State, allocator, ctx, start)).?;
+    const best, var graph: Map(State, Set(State)) = (try aoc.shortestPath(State, allocator, ctx, start, true)).?;
+    defer graph.deinit();
+    defer {
+        var it = graph.valueIterator();
+        while (it.next()) |set| set.deinit();
+    }
+
+    var onPath = Set(Pos).init(allocator);
+    defer onPath.deinit();
+
+    {
+        var dfs = List(State).init(allocator);
+        defer dfs.deinit();
+
+        inline for ([_]Direction{ .n, .s, .e, .w }) |dir| {
+            try dfs.append(State{
+                .pos = goal,
+                .dir = dir,
+            });
+        }
+
+        while (dfs.popOrNull()) |state| {
+            try onPath.put(state.pos);
+
+            var prevs = graph.get(state) orelse continue;
+            var it = prevs.iterator();
+            while (it.next()) |next| {
+                try dfs.append(next.*);
+            }
+        }
+    }
+
+    return .{ best, onPath.count() };
 }
 
 const Pos = struct {