aoc-2020/03/lib.ml

(* open Base *)
open Core

module PairCmp = struct
  module T = struct
    type t = int * int
    let compare x y = Tuple2.compare ~cmp1:Int.compare ~cmp2:Int.compare x y
    let sexp_of_t = Tuple2.sexp_of_t Int.sexp_of_t Int.sexp_of_t
    let t_of_sexp = Tuple2.t_of_sexp Int.t_of_sexp Int.t_of_sexp
  end
  include T
  include Comparable.Make(T)
end

let all_slopes = [
    (1, 1);
    (3, 1);
    (5, 1);
    (7, 1);
    (1, 2);
  ]

type map = {
    height: int;
    width: int;
    content: (PairCmp.t, bool, PairCmp.comparator_witness) Map.t;
  }

let parse_map str =
  let lines = String.split str ~on:'\n' in
  let assoc =
    List.mapi lines ~f:(fun row line ->
        List.mapi (String.to_list line) ~f:(fun col value ->
            ((col, row), phys_equal value '#'))) in
  {
    height = List.length assoc;
    width = List.length (List.nth_exn assoc 0);
    content = Map.of_alist_exn (module PairCmp) (List.join assoc);
  }

let count_trees_in_slope map (dx, dy) =
  let rec helper x y =
    if y >= map.height then 0
    else if x >= map.width then helper (x - map.width) y
    else if Map.find_exn map.content (x, y) then 1 + helper (x + dx) (y + dy)
    else helper (x + dx) (y + dy)
  in helper 0 0