aoc2020/day1.ml

module IntSet = Set_ext.Make(Int)

let read_ints =
  Bracket.infile_lines ~line:int_of_string ~of_seq:IntSet.of_seq


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' n input =
  match find ~n (read_ints input) with
  | Some xs -> Misc.print_prod xs
  | None    -> raise Not_found

let main = Misc.main 1 [|main' 2; main' 3|]


let%expect_test "part 1" =
  main' 2 "../../data/day1"; (* is this always the right path??? *)
  [%expect{| 534 * 1486 = 793524 |}]

let%expect_test "part 2" =
  main' 3 "../../data/day1";
  [%expect{| 71 * 686 * 1263 = 61515678 |}]