86 lines
2.9 KiB
OCaml
86 lines
2.9 KiB
OCaml
module IntSet = Set.Make (Int)
|
|
module IntHashtbl = Hashtbl.Make (Int)
|
|
|
|
(** [add_rule a b m] adds the rule that [a] must appear before [b] to the rule
|
|
map [m]. Returns the updated map [m] *)
|
|
let add_rule a b m =
|
|
match IntHashtbl.find_opt m a with
|
|
| None ->
|
|
IntHashtbl.add m a (IntSet.singleton b);
|
|
m
|
|
| Some s ->
|
|
IntHashtbl.replace m a (IntSet.add b s);
|
|
m
|
|
|
|
(** [find_rule a b m] returns [true] if the rule map [m] says that [a] should
|
|
appear before [b]. *)
|
|
let find_rule a b m =
|
|
match IntHashtbl.find_opt m a with
|
|
| Some s -> ( match IntSet.find_opt b s with Some _ -> true | None -> false)
|
|
| None -> false
|
|
|
|
(** [compare m a b] is a total ordering on pages in the rule map [m]. Returns
|
|
[-1] if [a] should appear before [b], [0] if [a = b], and [1] if [b] should
|
|
appear before [a]. *)
|
|
let compare m a b =
|
|
if a = b then 0
|
|
else if find_rule a b m then -1
|
|
else if find_rule b a m then 1
|
|
else failwith "compare"
|
|
|
|
(** [sort m pages] Sorts [pages] into the order required by the rule map [m]. *)
|
|
let sort m = List.sort (compare m)
|
|
|
|
(** [is_page_order_valid m pages] returns [true] iff the page ordering given in
|
|
[pages] is valid under the rule map [m]. *)
|
|
let rec is_page_order_valid m pages =
|
|
let rec impl h = function
|
|
| h' :: t -> find_rule h h' m && impl h t
|
|
| [] -> true
|
|
in
|
|
match pages with h :: t -> impl h t && is_page_order_valid m t | [] -> true
|
|
|
|
(** [parse_rules lst] parses the rules in the list [lst] stopping when
|
|
encountering an empty line. Returns a pair [(rule_map, tail)].
|
|
[tail] starts the line after the empty line. *)
|
|
let parse_rules =
|
|
let re = Str.regexp_string "|" in
|
|
let rec impl acc = function
|
|
| "" :: t -> (acc, t)
|
|
| [] -> failwith "parse_rules.impl"
|
|
| h :: t -> (
|
|
match List.map int_of_string (Str.split re h) with
|
|
| [ a; b ] -> impl (add_rule a b acc) t
|
|
| _ -> failwith "parse_rules.impl")
|
|
in
|
|
impl (IntHashtbl.create 17)
|
|
|
|
(** [parse_page_orders lst] parses a list of page orders. *)
|
|
let parse_page_orders =
|
|
let re = Str.regexp_string "," in
|
|
let rec impl acc = function
|
|
| [] -> acc
|
|
| h :: t -> impl (List.map int_of_string (Str.split re h) :: acc) t
|
|
in
|
|
impl []
|
|
|
|
(** Read a rule map and list of pages from the file [fname]. *)
|
|
let read_file fname =
|
|
let lst = Aoc.strings_of_file fname in
|
|
let rule_map, t = parse_rules lst in
|
|
let page_orders = parse_page_orders t in
|
|
(rule_map, page_orders)
|
|
|
|
let middle_elt lst = List.nth lst (List.length lst / 2)
|
|
|
|
let part1 (rule_map, page_orders) =
|
|
List.filter (is_page_order_valid rule_map) page_orders
|
|
|> List.map middle_elt |> List.fold_left ( + ) 0
|
|
|
|
let part2 (rule_map, page_orders) =
|
|
List.filter (fun lst -> not (is_page_order_valid rule_map lst)) page_orders
|
|
|> List.map (sort rule_map)
|
|
|> List.map middle_elt |> List.fold_left ( + ) 0
|
|
|
|
let _ = Aoc.main read_file [ (string_of_int, part1); (string_of_int, part2) ]
|