aoc-2022-ocaml/src/utils.ml

132 lines
3.1 KiB
OCaml

open Containers
let lines_of_input day =
let base_path =
match Inputs.Sites.inputs with
| path :: _ -> path
| _ -> failwith "cant find path containing inputs"
in
let path = Printf.sprintf "%s/day%d.tt" base_path day in
IO.(with_in path read_lines_l)
;;
(** [f] in [memo f] should be a function that makes recursive calls
via the memoized function passed as its first argument. ie:
{[
let fib_m =
memo (fun self x ->
match x with
| 0 | 1 -> 1
| x -> self (x - 1) + self (x - 2))
;;
]} *)
let memo f =
let open Hashtbl in
let cache = create 100000 in
let rec f_mem k =
try find cache k with
| Not_found ->
let v = f f_mem k in
add cache k v;
v
in
f_mem
;;
let%expect_test "fib memoized" =
let fib_m =
memo (fun fib_m x ->
match x with
| 0 | 1 -> 1
| x -> fib_m (x - 1) + fib_m (x - 2))
in
Printf.printf "%i\n" @@ fib_m 100;
[%expect {| 1298777728820984005 |}]
;;
let djikstra ~eq ~get_info ~set_info ~get_weight ~get_neighbors ~goal ~start_nodes =
let seen = ref [] in
let rec djikstra' c_node =
let c_dist, _ = get_info c_node in
set_info c_node (c_dist, true);
get_neighbors c_node
|> List.iter (fun neighbor ->
let n_dist, n_visited = get_info neighbor in
if not n_visited
then (
let new_dist = c_dist + get_weight c_node neighbor in
if new_dist < n_dist then set_info neighbor (new_dist, n_visited);
seen := List.add_nodup ~eq neighbor !seen));
seen
:= List.fast_sort
(fun a b ->
let a_dist, _ = get_info a in
let b_dist, _ = get_info b in
Int.compare a_dist b_dist)
!seen;
match goal, !seen with
| Some goal_node, next :: _ when eq goal_node next ->
let dist, _ = get_info next in
Some dist
| None, [] -> Some 0
| _, next :: rest ->
seen := rest;
djikstra' next
| Some _, [] -> None
in
List.iter (fun node -> set_info node (0, false)) start_nodes;
match start_nodes with
| first :: rest ->
seen := rest;
djikstra' first
| _ -> failwith "start_nodes cannot be empty"
;;
let djikstra_hash
(type k)
?(hash = Hashtbl.hash)
?(eq = Stdlib.( = ))
?(slots = 100)
~get_weight
~get_neighbors
~start_nodes
=
let module Tbl =
CCHashtbl.Make (struct
type t = k
let equal = eq
let hash = hash
end)
in
let table = Tbl.create slots in
let get_info node = Tbl.get_or ~default:(Int.max_int, false) table node in
let set_info node info = Tbl.replace table node info in
let dist_of node = Option.(Tbl.get table node >|= fun (dist, _) -> dist) in
let _ =
djikstra ~eq ~get_info ~set_info ~get_weight ~get_neighbors ~goal:None ~start_nodes
in
dist_of
;;
module Parse = struct
include Angstrom
let sign = option 1 (char '-' >>= fun _ -> return (-1))
let digits =
take_while1 (function
| '0' .. '9' -> true
| _ -> false)
>>| int_of_string
;;
let int = map2 sign digits ~f:( * )
let not_int =
skip_while (function
| '0' .. '9' | '-' -> false
| _ -> true)
;;
end