add expect tests & refactor

day1.ml

@@ -1,36 +1,53 @@
-module IntSet = Set.Make(Int)
+module IntSet = struct
+  include Set.Make(Int)
+
+  let exists_opt (type a) (f: int -> a option) (set: t): a option =
+    let exception Found of a in
+    let f' x = Option.iter (fun x -> raise (Found x)) (f x) in
+    match iter f' set with
+    | exception Found x -> Some x
+    | _                 -> None
+end
 
 let read_ints =
   Bracket.infile_lines ~line:int_of_string ~of_seq:IntSet.of_seq
 
-let find_with target ints =
-  let open IntSet in
-  let p x = mem (target - x) ints in
-  let res = min_elt (filter p ints) in
-  res, target - res
 
-let find = find_with 2020
+let rec find ?(target=2020) ~n set =
+  if n < 1 then
+    raise (Invalid_argument "n must be at least 1")
+  else if n = 1 then
+    if IntSet.mem target set then Some [target] else None
+  else
+    let joined x =
+      Option.map (List.cons x) (find ~target:(target - x) ~n:(n - 1) set) in
+    IntSet.exists_opt joined set
 
-let main_part_1 input =
-  let ints = read_ints input in
-  let x, y = find ints in
-  Printf.printf "%d * %d = %d\n" x y (x * y)
+let print_result xs =
+  let res = List.fold_left ( * ) 1 xs in
+  let rec print_prod fmt = function
+    | []    -> Format.pp_print_text fmt "{}"
+    | [x]   -> Format.fprintf fmt "%d" x
+    | x::xs -> Format.fprintf fmt "%d * %a" x print_prod xs in
+  Format.printf "%a = %d\n" print_prod xs res
 
-let main_part_2 input =
-  let ints = read_ints input in
-  let exception Found of int * int * int in
-  let go x =
-    match find_with (2020 - x) ints with
-    | exception Not_found -> ()
-    | y, z -> raise (Found (x, y, z))
-  in
-  match IntSet.iter go ints with
-  | exception Found (x, y, z) ->
-      Printf.printf "%d * %d * %d = %d\n" x y z (x * y * z)
-  | _ -> raise Not_found
 
-let mains = [|main_part_1; main_part_2|]
+let main' n input =
+  match find ~n (read_ints input) with
+  | Some xs -> print_result xs
+  | None    -> raise Not_found
+
+let mains = [|main' 2; main' 3|]
 
 let main = function
   | [part; input] -> mains.(int_of_string part - 1) input
-  | _ -> Usage.exit "1 <part> <infile>"
+  | _             -> Usage.exit "1 <part> <infile>"
+
+
+let%expect_test "part 1" =
+  mains.(0) "../../data/day1"; (* is this always the right path??? *)
+  [%expect{| 534 * 1486 = 793524 |}]
+
+let%expect_test "part 2" =
+  mains.(1) "../../data/day1";
+  [%expect{| 71 * 686 * 1263 = 61515678 |}]