108 lines
3.9 KiB
OCaml
108 lines
3.9 KiB
OCaml
module StringMap = Map.Make (String)
|
|
module StringSet = Set.Make (String)
|
|
|
|
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 rec find_second_member acc connections visited a candidates =
|
|
let rec impl acc set = function
|
|
| [] -> acc
|
|
| c :: t -> impl (StringSet.add c set :: acc) set t
|
|
in
|
|
match StringSet.choose_opt candidates with
|
|
| None -> acc
|
|
| Some h ->
|
|
let candidates = StringSet.remove h candidates in
|
|
if StringSet.mem h visited then
|
|
find_second_member acc connections visited a candidates
|
|
else
|
|
let visited = StringSet.add h visited in
|
|
let anh = StringSet.inter (StringMap.find h connections) candidates in
|
|
let acc =
|
|
impl acc (StringSet.of_list [ a; h ]) (StringSet.to_list anh)
|
|
in
|
|
find_second_member acc connections visited a candidates
|
|
|
|
let rec find_rings acc visited connections = function
|
|
| [] -> 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
|
|
(StringMap.find h connections)
|
|
in
|
|
find_rings acc visited connections t
|
|
|
|
(** [starts_with_t set] returns true if any member of [set] starts with the
|
|
letter ['t']. *)
|
|
let starts_with_t set = StringSet.exists (fun x -> x.[0] = 't') set
|
|
|
|
let part1 connections =
|
|
StringMap.to_list connections
|
|
|> List.map fst
|
|
|> find_rings [] StringSet.empty connections
|
|
|> List.filter starts_with_t |> List.length |> string_of_int
|
|
|
|
(** [search_candidate max_lst connections current candidates] searches for the
|
|
longest star network that contains all the computers in [current] and maybe
|
|
some of the elements of [candidates].
|
|
|
|
[max_lst] is the biggest network found so far. [connections] is the map of
|
|
computers to direct connections.
|
|
|
|
The initial call should look something like:
|
|
|
|
[search_candidate [] connections k (StringSet.to_list (StringMap.find k
|
|
connections))] *)
|
|
let rec search_candidate max_lst connections current candidates =
|
|
match candidates with
|
|
| [] -> if List.length current > List.length max_lst then current else max_lst
|
|
| h :: t ->
|
|
let map = StringMap.find h connections in
|
|
let current' = h :: current in
|
|
let candidates' = List.filter (Fun.flip StringSet.mem map) candidates in
|
|
let max_lst = search_candidate max_lst connections current' candidates' in
|
|
search_candidate max_lst connections current t
|
|
|
|
(** [find_max_set connections] returns a list containing the largest number of
|
|
computers in a star network (that is for every pair of elements in the list
|
|
there is a connection between them).
|
|
|
|
[connections] is the map of connections keyed by computer with the value
|
|
being a set of all direct connections. [connections] must be bi-directional
|
|
that is if: [StringSet.mem a (StringMap.find b connections)] then
|
|
[StringSet.mem b (StringMap.find a connections)]. Note that
|
|
[StringSet.mem a (StringMap.find a connections)] must return [false]. *)
|
|
let find_max_set connections =
|
|
let rec impl max_lst = function
|
|
| [] -> max_lst
|
|
| (k, v) :: t ->
|
|
impl
|
|
(search_candidate max_lst connections [ k ] (StringSet.to_list v))
|
|
t
|
|
in
|
|
impl [] (StringMap.to_list connections)
|
|
|
|
let part2 connections =
|
|
find_max_set connections |> List.sort compare |> String.concat ","
|
|
|
|
let _ = Aoc.main load_file [ (Fun.id, part1); (Fun.id, part2) ]
|