54 lines
1.5 KiB
OCaml
54 lines
1.5 KiB
OCaml
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 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 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' 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>"
|
|
|
|
|
|
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 |}]
|