Day 5 pt2, getting fiddly already!

1 files changed, 241 insertions, 32 deletions

diff --git a/day5.zig b/day5.zig
index 572102d..bda0c74 100644
--- a/day5.zig
+++ b/day5.zig
@@ -2,35 +2,17 @@ const std = @import("std");
 
 pub fn main() !void {
     try std.fmt.format(std.io.getStdOut().writer(), "day5 pt1: {}\n", .{try solve_pt1(std.heap.page_allocator, puzzle_input)});
+    try std.fmt.format(std.io.getStdOut().writer(), "day5 pt2: {}\n", .{try solve_pt2(std.heap.page_allocator, puzzle_input)});
 }
 
 fn solve_pt1(a: std.mem.Allocator, input: []const u8) !u64 {
-    var groups = std.mem.split(u8, input, "\n\n");
-    const seeds_str = groups.first();
-    var seedToks = std.mem.tokenize(u8, seeds_str, " ");
-    _ = seedToks.next() orelse return error.NoSeeds; // ignore "seeds:"
-    var seeds = std.ArrayList(u64).init(a);
-    defer seeds.deinit();
-    while (seedToks.next()) |seedTok| {
-        try seeds.append(try std.fmt.parseInt(u64, seedTok, 10));
-    }
-
-    // now parse the maps...
-    var maps = std.ArrayList(Map).init(a);
-    defer {
-        for (maps.items) |*map| map.deinit();
-        maps.deinit();
-    }
-    while (groups.next()) |group| {
-        const m = try Map.parse(a, group);
-        try maps.append(m);
-    }
-
+    var puzzle = try Puzzle.parse(a, input);
+    defer puzzle.deinit(a);
     // and now check the seeds
     var min: u64 = std.math.maxInt(u64);
-    for (seeds.items) |seed| {
+    for (puzzle.seeds.items) |seed| {
         var location = seed;
-        for (maps.items) |map| {
+        for (puzzle.maps.items) |map| {
             location = map.apply(location);
         }
         min = @min(min, location);
@@ -38,34 +20,242 @@ fn solve_pt1(a: std.mem.Allocator, input: []const u8) !u64 {
     return min;
 }
 
+// Naïvely trying to iterate results in too much calculation, so we have to be a bit smarter
+// and handle only the ends of contiguous ranges. The Map.applyRange function will take a range 
+// and apply the map, splitting it into sub-ranges if necessary e.g. because a mapping only partially
+// overlaps a range.
+fn solve_pt2(a: std.mem.Allocator, input: []const u8) !u64 {
+    var puzzle = try Puzzle.parse(a, input);
+    defer puzzle.deinit(a);
+
+    // Parse all the ranges of seeds
+    var ranges = std.ArrayList(Range).init(a);
+    defer ranges.deinit();
+    for (0..(puzzle.seeds.items.len / 2)) |i| {
+        const start = puzzle.seeds.items[2 * i];
+        const len = puzzle.seeds.items[(2 * i) + 1];
+        try ranges.append(.{
+            .start = start,
+            .end = start + len,
+        });
+    }
+
+    // Go through the maps, map the ranges each time.
+    for (puzzle.maps.items) |map| {
+        const sr = try ranges.toOwnedSlice();
+        for (sr) |range| {
+            try ranges.appendSlice(try map.applyRange(a, range));
+        }
+        a.free(sr);
+    }
+
+    // Give me the start of the range with the lowest start
+    const fr = try ranges.toOwnedSlice();
+    defer a.free(fr);
+    std.mem.sort(Range, fr, {}, Range.cmp);
+    return fr[0].start;
+}
+
+const Puzzle = struct {
+    seeds: std.ArrayList(u64),
+    maps: std.ArrayList(Map),
+
+    fn parse(a: std.mem.Allocator, input: []const u8) !Puzzle {
+        var groups = std.mem.split(u8, input, "\n\n");
+        const seeds_str = groups.first();
+        var seedToks = std.mem.tokenize(u8, seeds_str, " ");
+        _ = seedToks.next() orelse return error.NoSeeds; // ignore "seeds:"
+        var seeds = std.ArrayList(u64).init(a);
+        errdefer seeds.deinit();
+        while (seedToks.next()) |seedTok| {
+            try seeds.append(try std.fmt.parseInt(u64, seedTok, 10));
+        }
+
+        // now parse the maps...
+        var maps = std.ArrayList(Map).init(a);
+        errdefer {
+            for (maps.items) |*map| map.deinit(a);
+            maps.deinit();
+        }
+        while (groups.next()) |group| {
+            const m = try Map.parse(a, group);
+            try maps.append(m);
+        }
+        return .{
+            .seeds = seeds,
+            .maps = maps,
+        };
+    }
+    fn deinit(self: *Puzzle, a: std.mem.Allocator) void {
+        self.seeds.deinit();
+        for (self.maps.items) |*map| map.deinit(a);
+        self.maps.deinit();
+    }
+};
+
 const Map = struct {
-    entries: std.ArrayList(MapEntry),
+    entries: []MapEntry, // sorted by start
 
     fn parse(a: std.mem.Allocator, input: []const u8) !Map {
         var entries = std.ArrayList(MapEntry).init(a);
-        errdefer entries.deinit();
+        defer entries.deinit();
         var toks = std.mem.split(u8, input, "\n");
         _ = toks.first(); // discard the title.
         while (toks.next()) |tok| {
             const me = try MapEntry.parse(tok);
             try entries.append(me);
         }
+        const ee = try entries.toOwnedSlice();
+        std.mem.sort(MapEntry, ee, {}, Map.cmp);
         return .{
-            .entries = entries,
+            .entries = ee,
         };
     }
 
-    fn deinit(self: *Map) void {
-        self.entries.deinit();
+    fn cmp(x: void, lhs: MapEntry, rhs: MapEntry) bool {
+        _ = x;
+        return lhs.source_start < rhs.source_start;
+    }
+
+    fn deinit(self: *Map, a: std.mem.Allocator) void {
+        a.free(self.entries);
     }
 
     fn apply(self: Map, in: u64) u64 {
-        for (self.entries.items) |entry| {
+        for (self.entries) |entry| {
             if (entry.apply(in)) |out| return out;
         } else {
             return in;
         }
     }
+
+    fn applyRange(self: Map, a: std.mem.Allocator, r: Range) ![]Range {
+        var res = std.ArrayList(Range).init(a);
+        defer res.deinit();
+        var r2 = r; // mutable range for bits we have left... 
+        for (self.entries) |me| {
+            if (me.source_end < r2.start) {
+                // ignore this map entry it's out of bounds
+            } else if (me.source_start >= r2.end) {
+                // ignore this map entry it's out of bounds
+            } else if (me.source_start <= r2.start and me.source_end >= r2.end) {
+                // whole range contained by map entry
+                try res.append(.{
+                    .start = me.apply(r2.start) orelse @panic("boo"),
+                    .end = (me.apply(r2.end-1) orelse @panic("bar"))+1,
+                });
+                const x = try res.toOwnedSlice();
+                std.mem.sort(Range, x, {}, Range.cmp);
+                return x;
+            } else if (me.source_start < r2.start and me.source_end > r2.start and me.source_end < r2.end) {
+                // overlap start of range
+                try res.append(.{
+                    .start = me.apply(r2.start) orelse @panic("glub"),
+                    .end = ( me.apply(me.source_end-1) orelse @panic("plub")) + 1,
+                });
+                const ne = r2.end;
+                r2 = .{
+                    .start = me.source_end,
+                    .end = ne,
+                };
+            } else if (me.source_start >= r2.start and me.source_end > r2.end) {
+                // overlap end of range
+                try res.append(.{
+                    .start = r2.start,
+                    .end = me.source_start,    
+                });
+                try res.append(.{
+                    .start = me.apply(me.source_start) orelse @panic("ploop"),
+                    .end = (me.apply(r2.end-1) orelse @panic("gloop")) + 1,
+                });
+                const x = try res.toOwnedSlice();
+                std.mem.sort(Range, x, {}, Range.cmp);
+                return x;
+            } else if (me.source_start >= r2.start and me.source_end < r2.end) {
+                // map contained by range, 3 way split... 
+                // overlap end of range
+                try res.append(.{
+                    .start = r2.start,
+                    .end = me.source_start,    
+                });
+                try res.append(.{
+                    .start = me.apply(me.source_start) orelse @panic("ploop"),
+                    .end = (me.apply(me.source_end-1) orelse @panic("gloop")) + 1,
+                });
+                const ne = r2.end;
+                r2 = .{
+                    .start = me.source_end,
+                    .end = ne,
+                };
+            }
+        } else {
+            // we had an unmapped section, retain it (equivalent to identity map)
+            try res.append(r2);
+        }
+        const x = try res.toOwnedSlice();
+        std.mem.sort(Range, x, {}, Range.cmp);
+        return x;
+    }
+
+    test "applyRange no overlap" {
+        const a = std.testing.allocator;
+        var m = try Map.parse(a, 
+            \\foobar!
+            \\3 10 2
+        );
+        defer m.deinit(a);
+        const res = try m.applyRange(a, .{.start = 8, .end = 10});
+        defer a.free(res);
+        try std.testing.expectEqualSlices(Range, &[_]Range{.{.start = 8, .end = 10}}, res);
+
+        const res2 = try m.applyRange(a, .{.start = 12, .end = 13});
+        defer a.free(res2);
+        try std.testing.expectEqualSlices(Range, &[_]Range{.{.start = 12, .end = 13}}, res2); 
+    }
+    test "applyRange overlap start" {
+        const a = std.testing.allocator;
+        var m = try Map.parse(a, 
+            \\foobar!
+            \\5 20 10
+        );
+        defer m.deinit(a);
+        const res = try m.applyRange(a, .{.start = 25, .end  = 35});
+        defer a.free(res);
+        try std.testing.expectEqualSlices(Range, &[_]Range{.{.start = 10, .end = 15}, .{.start = 30, .end = 35}}, res);
+    }
+    test "applyRange overlap end" {
+        const a = std.testing.allocator;
+        var m = try Map.parse(a, 
+            \\foobar!
+            \\5 20 10
+        );
+        defer m.deinit(a);
+        const res = try m.applyRange(a, .{.start = 15, .end  = 25});
+        defer a.free(res);
+        try std.testing.expectEqualSlices(Range, &[_]Range{.{.start = 5, .end = 10}, .{.start = 15, .end = 20}}, res);
+    }
+    test "applyRange map contains whole range" {
+        const a = std.testing.allocator;
+        var m = try Map.parse(a, 
+            \\foobar!
+            \\5 20 10
+        );
+        defer m.deinit(a);
+        const res = try m.applyRange(a, .{.start = 20, .end  = 30});
+        defer a.free(res);
+        try std.testing.expectEqualSlices(Range, &[_]Range{.{.start = 5, .end = 15}}, res);
+    }
+    test "applyRange range contains whole map" {
+        const a = std.testing.allocator;
+        var m = try Map.parse(a, 
+            \\foobar!
+            \\5 20 10
+        );
+        defer m.deinit(a);
+        const res = try m.applyRange(a, .{.start = 15, .end  = 35});
+        defer a.free(res);
+        try std.testing.expectEqualSlices(Range, &[_]Range{.{.start = 5, .end = 15}, .{.start = 15, .end = 20}, .{.start = 30, .end = 35}}, res);
+    }
 };
 
 const MapEntry = struct {
@@ -73,12 +263,20 @@ const MapEntry = struct {
     source_start: u64,
     range_len: u64,
 
+    source_end: u64, // exclusive bound
+    dest_end: u64,   // exclusive bound
+
     fn parse(input: []const u8) !MapEntry {
         var toks = std.mem.tokenize(u8, input, " ");
+        const dest_start = try std.fmt.parseInt(u64, (toks.next() orelse return error.NoDest), 10);
+        const source_start = try std.fmt.parseInt(u64, (toks.next() orelse return error.NoSource), 10);
+        const range_len = try std.fmt.parseInt(u64, (toks.next() orelse return error.NoRange), 10);
         return .{
-            .dest_start = try std.fmt.parseInt(u64, (toks.next() orelse return error.NoDest), 10),
-            .source_start = try std.fmt.parseInt(u64, (toks.next() orelse return error.NoSource), 10),
-            .range_len = try std.fmt.parseInt(u64, (toks.next() orelse return error.NoRange), 10),
+            .dest_start = dest_start,
+            .source_start = source_start,
+            .range_len = range_len,
+            .source_end = source_start + range_len,
+            .dest_end = dest_start + range_len,
         };
     }
 
@@ -91,9 +289,20 @@ const MapEntry = struct {
     }
 };
 
+const Range = struct {
+    start: u64,
+    end: u64,
+    fn cmp(_:void, lhs: Range, rhs: Range) bool {
+        return lhs.start < rhs.start;
+    }
+};
+
 test "pt1" {
     try std.testing.expectEqual(@as(u64, 35), try solve_pt1(std.testing.allocator, test_input));
 }
+test "pt2" {
+    try std.testing.expectEqual(@as(u64, 46), try solve_pt2(std.testing.allocator, test_input));
+}
 
 const test_input =
     \\seeds: 79 14 55 13