fix day 15 for example inputs

part 1 didn't work for the example input :(

src/day15.ml

@@ -1,5 +1,23 @@
 open Containers
 
+let example_lines =
+  [ "Sensor at x=2, y=18: closest beacon is at x=-2, y=15"
+  ; "Sensor at x=9, y=16: closest beacon is at x=10, y=16"
+  ; "Sensor at x=13, y=2: closest beacon is at x=15, y=3"
+  ; "Sensor at x=12, y=14: closest beacon is at x=10, y=16"
+  ; "Sensor at x=10, y=20: closest beacon is at x=10, y=16"
+  ; "Sensor at x=14, y=17: closest beacon is at x=10, y=16"
+  ; "Sensor at x=8, y=7: closest beacon is at x=2, y=10"
+  ; "Sensor at x=2, y=0: closest beacon is at x=2, y=10"
+  ; "Sensor at x=0, y=11: closest beacon is at x=2, y=10"
+  ; "Sensor at x=20, y=14: closest beacon is at x=25, y=17"
+  ; "Sensor at x=17, y=20: closest beacon is at x=21, y=22"
+  ; "Sensor at x=16, y=7: closest beacon is at x=15, y=3"
+  ; "Sensor at x=14, y=3: closest beacon is at x=15, y=3"
+  ; "Sensor at x=20, y=1: closest beacon is at x=15, y=3"
+  ]
+;;
+
 type sensor =
   { pos : Vec2.t
   ; beacon : Vec2.t
@@ -51,24 +69,23 @@ let could_be_beacon sensors ignore_known pos =
       n1 = y + bx + dist + 1
       n2 = y + bx - dist - 1
   ]}*)
-let get_consts sensors =
+let get_consts sensors extra =
   sensors
   |> Iter.fold
        (fun (p_slope_consts, n_slope_consts) { pos; beacon = _; dist } ->
-         let p1 = pos.y - pos.x + dist + 1 in
+         let p1 = pos.y - pos.x + dist + extra in
-         let p2 = pos.y - pos.x - dist - 1 in
+         let p2 = pos.y - pos.x - dist - extra in
-         let n1 = pos.y + pos.x + dist + 1 in
+         let n1 = pos.y + pos.x + dist + extra in
-         let n2 = pos.y + pos.x - dist - 1 in
+         let n2 = pos.y + pos.x - dist - extra in
          Iter.(
            ( append p_slope_consts (of_list [ p1; p2 ])
            , append n_slope_consts (of_list [ n1; n2 ]) )))
        (Iter.empty, Iter.empty)
 ;;
 
-let%expect_test "Day 15.1" =
+let solve_part_1 lines scan_y =
-  let sensors = Utils.lines_of_input 15 |> Iter.of_list |> Iter.map parse_sensor in
+  let sensors = lines |> Iter.of_list |> Iter.map parse_sensor in
-  let scan_y = 2000000 in
+  let p_slope_consts, n_slope_consts = get_consts sensors 0 in
-  let p_slope_consts, n_slope_consts = get_consts sensors in
   (*{[
       for a known y to intersect with a positive slope boundary:
         y = x + p
@@ -84,31 +101,41 @@ let%expect_test "Day 15.1" =
       (Iter.map (fun n -> n - scan_y) n_slope_consts)
     |> Iter.sort_uniq ~cmp:Int.compare
   in
-  (* iterate thru the intersections by ascending x.
-     if a point in the region between two intersections can be a beacon,
-     then every point in the region can be. *)
   let rec count_not_beacons count xs =
     match xs with
-    | lx :: hx :: rest ->
+    | x :: next :: rest ->
-      let count =
+      let intersection =
-        if could_be_beacon sensors false @@ Vec2.of_tuple (lx + 1, scan_y)
+        if could_be_beacon sensors false @@ Vec2.of_tuple (x, scan_y) then 0 else 1
-        then count
-        else count + hx - lx
       in
-      count_not_beacons count @@ (hx :: rest)
+      (* if there is a region between this intersection and the next, the value
+         of could_be_beacon is the same for each point within it. *)
+      let region =
+        if x + 1 < next && (could_be_beacon sensors false @@ Vec2.of_tuple (x + 1, scan_y))
+        then 0
+        else next - x - 1
+      in
+      count_not_beacons (count + intersection + region) (next :: rest)
     | _ -> count
   in
-  let possible_beacons = count_not_beacons 0 @@ Iter.to_list intersection_xs in
+  count_not_beacons 0 @@ Iter.to_list intersection_xs
+;;
+
+let%expect_test "Day 15.1 example" =
+  let possible_beacons = solve_part_1 example_lines 10 in
+  Printf.printf "%i\n" @@ possible_beacons;
+  [%expect {| 26 |}]
+;;
+
+let%expect_test "Day 15.1" =
+  let possible_beacons = solve_part_1 (Utils.lines_of_input 15) 2000000 in
   Printf.printf "%i\n" @@ possible_beacons;
   [%expect {| 5525990 |}]
 ;;
 
-let%expect_test "Day 15.2" =
+let solve_part_2 lines range mult =
-  let sensors = Utils.lines_of_input 15 |> Iter.of_list |> Iter.map parse_sensor in
+  let sensors = lines |> Iter.of_list |> Iter.map parse_sensor in
-  let mult = 4000000 in
-  let range = mult in
   let in_bounds = Vec2.(in_bounds origin @@ of_tuple (range, range)) in
-  let p_slope_consts, n_slope_consts = get_consts sensors in
+  let p_slope_consts, n_slope_consts = get_consts sensors 1 in
   (*{[
       boundaries only itersect if they have opposite slope.
       to find the intersection, set equations to be equal and solve.
@@ -132,6 +159,15 @@ let%expect_test "Day 15.2" =
          None
     |> Option.get_exn_or "no distress beacon found :("
   in
-  Printf.printf "%i\n" ((mult * point.x) + point.y);
+  (mult * point.x) + point.y
+;;
+
+let%expect_test "Day 15.2 example" =
+  Printf.printf "%i\n" @@ solve_part_2 example_lines 20 4000000;
+  [%expect {| 56000011 |}]
+;;
+
+let%expect_test "Day 15.2" =
+  Printf.printf "%i\n" @@ solve_part_2 (Utils.lines_of_input 15) 4000000 4000000;
   [%expect {| 11756174628223 |}]
 ;;