119 lines
3.6 KiB
OCaml
119 lines
3.6 KiB
OCaml
module StringMap = Map.Make (String)
|
|
module StringSet = Set.Make (String)
|
|
|
|
let[@warning "-32"] print_set set =
|
|
StringSet.to_list set |> List.sort compare |> String.concat ","
|
|
|> Printf.printf "%s\n"
|
|
|
|
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 rec search_candidate max_set connections current candidates =
|
|
match StringSet.choose_opt candidates with
|
|
| None ->
|
|
if StringSet.cardinal current > StringSet.cardinal max_set then current
|
|
else max_set
|
|
| Some h ->
|
|
let current' = StringSet.add h current in
|
|
let candidates' =
|
|
StringSet.inter candidates (StringMap.find h connections)
|
|
in
|
|
let max_set = search_candidate max_set connections current' candidates' in
|
|
search_candidate max_set connections current
|
|
(StringSet.remove h candidates)
|
|
|
|
let find_max_set connections =
|
|
let rec impl max_set = function
|
|
| [] -> max_set
|
|
| (k, v) :: t ->
|
|
impl
|
|
(search_candidate max_set connections
|
|
(StringSet.add k StringSet.empty)
|
|
v)
|
|
t
|
|
in
|
|
impl StringSet.empty (StringMap.to_list connections)
|
|
|
|
let part2 connections =
|
|
let max_set = find_max_set connections in
|
|
print_set max_set;
|
|
StringSet.cardinal max_set
|
|
|
|
let _ = Aoc.main load_file [ (string_of_int, part1); (string_of_int, part2) ]
|