day3.ml

open Core

type action = [ `Do of int * int | `Dont of int * int ]

let sum_muls ~pattern ~input : int =
  let re = Re.compile (Re.Perl.re pattern) in

  Re.all re input
  |> List.fold ~init:0 ~f:(fun acc group ->
         let num1 = Re.Group.get group 1 |> Int.of_string in
         let num2 = Re.Group.get group 2 |> Int.of_string in
         acc + (num1 * num2))

let build_pairs (matches : action list) : action list =
  let rec aux acc enabled = function
    | [] -> List.rev acc
    | `Do (start, stop) :: rest when not enabled ->
        aux (`Do (start, stop) :: acc) true rest
    | `Dont (start, stop) :: rest when enabled ->
        aux (`Dont (start, stop) :: acc) false rest
    | _ :: rest -> aux acc enabled rest
  in
  aux [] false matches

let substr_input (match_list : action list) input =
  let rec aux acc = function
    | [] -> List.rev acc
    | `Do (_, stop) :: `Dont (start2, _) :: rest ->
        let sub = String.sub input ~pos:stop ~len:(start2 - stop) in
        aux (sub :: acc) rest
    | _ -> acc
  in
  aux [] match_list

let scan_positions re input : action list =
  Re.all re input
  |> List.filter_map ~f:(fun el ->
         let matched = Re.Group.get el 1 in
         let start_pos = Re.Group.start el 1 in
         let end_pos = Re.Group.stop el 1 in
         match matched with
         | "do()" -> Some (`Do (start_pos, end_pos))
         | "don't()" -> Some (`Dont (start_pos, end_pos))
         | _ -> None)

let part2 ~input ~mul_pattern =
  let do_pattern = "do\\(\\)" in
  let dont_pattern = "don\\'t\\(\\)" in
  let combined_pattern = "(" ^ do_pattern ^ "|" ^ dont_pattern ^ ")" in
  let new_input = "do()" ^ input in
  let re = Re.compile (Re.Perl.re combined_pattern) in
  let positions = scan_positions re new_input in
  let pairs = build_pairs positions in
  let substrings = substr_input pairs new_input in
  List.fold substrings ~init:0 ~f:(fun acc substring ->
      acc + sum_muls ~pattern:mul_pattern ~input:substring)

let () =
  let pattern = "mul\\((\\d{1,3}),(\\d{1,3})\\)" in
  let input = In_channel.read_all "day3.txt" in
  let sum_value = sum_muls ~pattern ~input in
  printf "answer: %d\n" sum_value;
  printf "answer part 2: %d\n" (part2 ~input ~mul_pattern:pattern)