day 3 pt2

1 files changed, 119 insertions, 11 deletions

diff --git a/day3.zig b/day3.zig
index b078f76..4fe95b5 100644
--- a/day3.zig
+++ b/day3.zig
@@ -4,7 +4,8 @@ const Setup = @import("common.zig").Setup;
 pub fn main() !void {
     // var setup = try Setup.get();
     // defer setup.deinit();
-    try std.fmt.format(std.io.getStdOut().writer(), "day 3: {}\n", .{try solve(std.heap.page_allocator, input2)});
+    try std.fmt.format(std.io.getStdOut().writer(), "day 3 pt1: {}\n", .{try solve(std.heap.page_allocator, puzzle_input)});
+    try std.fmt.format(std.io.getStdOut().writer(), "day 3 pt2: {}\n", .{try solve_pt2(std.heap.page_allocator, puzzle_input)});
 }
 
 const Grid = struct {
@@ -12,43 +13,72 @@ const Grid = struct {
     width: usize,
     height: usize,
 
-    fn init(input: []const u8) !Grid {
+    spans: std.ArrayList(u32), // numbers we read from the grid
+    span_ix: []usize, // grid mapping of index into spans list
+
+    fn init(a: std.mem.Allocator, input: []const u8) !Grid {
+        var spix = try a.alloc(usize, input.len);
+        @memset(spix, std.math.maxInt(usize));
+
+        errdefer a.free(spix);
         const height = std.mem.count(u8, input, "\n")+1;
         const width = std.mem.indexOfScalar(u8, input, '\n') orelse return error.NoNewlines;
         return .{
             .buf = input,
             .width = width,
             .height = height,
+            .spans = std.ArrayList(u32).init(a),
+            .span_ix = spix,
         };
     }
     fn get(self: Grid, x: usize, y: usize) ?u8 {
         if (x >= self.width) return null;
         if (y >= self.height) return null;
         // width +1 becuase the buffer still has newlines in it
-        return self.buf[@intCast((self.width+1) * y + x)];
+        return self.buf[((self.width+1) * y) + x];
+    }
+
+    fn get_span_ix(self: Grid, x: usize, y: usize) ?usize {
+        if (x >= self.width) return null;
+        if (y >= self.height) return null;
+        const res = self.span_ix[((self.width+1) * y) + x];
+        if (res == std.math.maxInt(usize)) return null;
+        return res;
+    }
+
+    fn set_span_ix(self: *Grid, x: usize, y: usize, span_ix: usize) void {
+        self.span_ix[((self.width+1) * y) + x] = span_ix;
+    }
+
+    fn deinit(self: *Grid, a: std.mem.Allocator) void {
+        a.free(self.span_ix);
     }
 };
 
 fn solve(a: std.mem.Allocator, input: []const u8) !u32 {
     var spanBuf = std.ArrayList(u8).init(a);
+    var spanStart: usize = 0; // easier than backtracking
     defer spanBuf.deinit();
-    const grid = try Grid.init(input);
+    var grid = try Grid.init(a, input);
+    defer grid.deinit(a);
     var sum: u32 = 0;
     for (0..grid.height) |y| {
         spanBuf.clearRetainingCapacity();
         for (0..grid.width) |x| {
             const ch = grid.get(x, y) orelse return error.IndexError;
             if (std.ascii.isDigit(ch)) {
+                if (spanBuf.items.len == 0) {
+                    spanStart = x;
+                }
                 try spanBuf.append(ch);
             } 
             // end-of-span, either hit a non-digit or end of line
             if ((!std.ascii.isDigit(ch) or (x==grid.width-1)) and spanBuf.items.len > 0) {
                 const num = try std.fmt.parseInt(u32, spanBuf.items, 10);
-                const num_len = spanBuf.items.len;
                 spanBuf.clearRetainingCapacity();
                 // now check all the surrounding positions including diagonals, for non-digits and points
                 const search_x_end = x;
-                const search_x_start = std.math.sub(usize, search_x_end, (num_len+1)) catch search_x_end-num_len;
+                const search_x_start = std.math.sub(usize, spanStart, 1) catch spanStart;
                 const is_part: bool = br: for (search_x_start..(search_x_end)+1) |x_p| {
                     const search_y_start = std.math.sub(usize, y, 1) catch y;
                     const search_y_end = y+1;
@@ -73,10 +103,90 @@ fn solve(a: std.mem.Allocator, input: []const u8) !u32 {
     return sum;
 }
 
+fn solve_pt2(a: std.mem.Allocator, input: []const u8) !u32 {
+    var spanBuf = std.ArrayList(u8).init(a);
+    var spanStart: usize = 0; // easier than backtracking
+    defer spanBuf.deinit();
+    var grid = try Grid.init(a, input);
+    defer grid.deinit(a);
+    var allSpans = std.ArrayList(u32).init(a);
+    defer allSpans.deinit();
+    
+    for (0..grid.height) |y| {
+        spanBuf.clearRetainingCapacity();
+        for (0..grid.width) |x| {
+            const ch = grid.get(x, y) orelse return error.IndexError;
+            if (std.ascii.isDigit(ch)) {
+                if (spanBuf.items.len == 0) {
+                    spanStart = x;
+                }
+                try spanBuf.append(ch);
+            } 
+            // end-of-span, either hit a non-digit or end of line
+            if ((!std.ascii.isDigit(ch) or (x==grid.width-1)) and spanBuf.items.len > 0) {
+                const num = try std.fmt.parseInt(u32, spanBuf.items, 10);
+                spanBuf.clearRetainingCapacity();
+                const spanIx = allSpans.items.len;
+                try allSpans.append(num);
+                const mark_x_start = spanStart;
+                const mark_x_end = if (std.ascii.isDigit(ch)) x else x-1;
+                for (mark_x_start..(mark_x_end+1)) |mark_x| {
+                    grid.set_span_ix(mark_x, y, spanIx);
+                }
+            }
+        }
+    }
 
+    // Second pass, find the gears
+    var sum: u32 = 0;
+    var spixes = std.AutoHashMap(usize, bool).init(a);
+    defer spixes.deinit();
+    for (0..grid.height) |y| {
+        for (0..grid.width) |x| {
+            const ch = grid.get(x, y) orelse @panic("borken");
+            if (ch == '*') {
+                std.log.warn("gear at x{}y{}", .{x, y});
+                const mx = std.math.sub(usize, x, 1) catch x;
+                const my = std.math.sub(usize, y, 1) catch y;
 
-test {
-    const input = 
+                spixes.clearRetainingCapacity();
+                for (mx..(x+2)) |x_p| {
+                    for (my..(y+2)) |y_p| {
+                        if (grid.get_span_ix(x_p, y_p)) |spix| {
+                            try spixes.put(spix, true);
+                        }
+                    }
+                }
+                var ki = spixes.keyIterator();
+                if (ki.next()) |first| {
+                    var m = allSpans.items[first.*];
+                    //std.log.warn("span val m{}", .{m});
+                    if (ki.next()) |second| {
+                        var m2 = allSpans.items[second.*];
+                        //std.log.warn("span val m2{}", .{m2});
+                        // _EXACTLY_ two gears!
+                        if (ki.next() == null) {
+                            sum += (m*m2);
+                        }
+                    }
+                }
+            }
+        }
+    }
+    return sum;
+}
+
+
+
+test "pt1" {
+    try std.testing.expectEqual(@as(u32, 4361), try solve(std.testing.allocator, test_input));
+}
+
+test "pt2" {
+    try std.testing.expectEqual(@as(u32, 467835), try solve_pt2(std.testing.allocator, test_input));
+}
+
+const test_input = 
     \\467..114..
     \\...*......
     \\..35..633.
@@ -88,10 +198,8 @@ test {
     \\...$.*....
     \\.664.598..
     ;
-    try std.testing.expectEqual(@as(u32, 4361), try solve(std.testing.allocator, input));
-}
 
-const input2 = 
+const puzzle_input = 
 \\.......497...........................858...923...128..................227..801........487.....664...........................................
 \\436........765..............140.......+....................859.............*.........+.................960........668.......................
 \\...*982...........=..........=....203......266.263...375*....=...402....691..-....................*..........575....................13......