92 lines
2.7 KiB
OCaml
92 lines
2.7 KiB
OCaml
module StringMap = Map.Make (String)
|
|
module StringSet = Set.Make (String)
|
|
|
|
let[@warning "-32"] print_set set =
|
|
StringSet.iter
|
|
(fun x ->
|
|
print_string x;
|
|
print_char ' ')
|
|
set;
|
|
print_newline ()
|
|
|
|
let[@warning "-32"] print_sets = List.iter print_set
|
|
|
|
let[@warning "-32"] print_map =
|
|
StringMap.iter (fun k v ->
|
|
Printf.printf "%s: " k;
|
|
print_set v)
|
|
|
|
let add_connection map (a, b) =
|
|
let update s = function
|
|
| None -> Some (StringSet.add s StringSet.empty)
|
|
| Some set -> Some (StringSet.add s set)
|
|
in
|
|
let map = StringMap.update a (update b) map in
|
|
let map = StringMap.update b (update a) map in
|
|
map
|
|
|
|
let make_pairs = function
|
|
| [ a; b ] -> (a, b)
|
|
| _ -> raise (invalid_arg "make_pairs")
|
|
|
|
let load_file fname =
|
|
Aoc.strings_of_file fname
|
|
|> List.map (String.split_on_char '-')
|
|
|> List.map make_pairs
|
|
|> List.fold_left add_connection StringMap.empty
|
|
|
|
let is_three_ring connections a _ candidate =
|
|
StringSet.mem candidate (StringMap.find a connections)
|
|
|
|
let rec find_third_member acc connections visited a b candidates =
|
|
match candidates with
|
|
| [] -> acc
|
|
| h :: t ->
|
|
if StringSet.mem h visited then
|
|
find_third_member acc connections visited a b t
|
|
else if not (is_three_ring connections a b h) then
|
|
find_third_member acc connections visited a b t
|
|
else
|
|
let acc = StringSet.of_list [ a; b; h ] :: acc in
|
|
find_third_member acc connections visited a b t
|
|
|
|
let rec find_second_member acc connections visited a candidates =
|
|
match candidates with
|
|
| [] -> acc
|
|
| h :: t ->
|
|
if StringSet.mem h visited then
|
|
find_second_member acc connections visited a t
|
|
else
|
|
let visited = StringSet.add h visited in
|
|
let acc =
|
|
find_third_member acc connections visited a h
|
|
(StringSet.to_list (StringMap.find h connections))
|
|
in
|
|
find_second_member acc connections visited a t
|
|
|
|
let rec find_rings acc visited connections computers =
|
|
match computers with
|
|
| [] -> acc
|
|
| h :: t ->
|
|
if StringSet.mem h visited then find_rings acc visited connections t
|
|
else
|
|
let visited = StringSet.add h visited in
|
|
let acc =
|
|
find_second_member acc connections visited h
|
|
(StringSet.to_list (StringMap.find h connections))
|
|
in
|
|
find_rings acc visited connections t
|
|
|
|
let starts_with_t set = StringSet.exists (fun x -> x.[0] = 't') set
|
|
|
|
let part1 connections =
|
|
let computers = StringMap.to_list connections |> List.map fst in
|
|
let rings =
|
|
find_rings [] StringSet.empty connections computers
|
|
|> List.filter starts_with_t
|
|
in
|
|
print_sets rings;
|
|
List.length rings
|
|
|
|
let _ = Aoc.main load_file [ (string_of_int, part1) ]
|