100 lines
3.2 KiB
OCaml
100 lines
3.2 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 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
|
|
string_of_int (List.length rings)
|
|
|
|
let rec search_candidate max_set connections current candidates =
|
|
match StringSet.choose_opt candidates with
|
|
| None ->
|
|
if List.length current > List.length max_set then current else max_set
|
|
| Some h ->
|
|
let current' = 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 [ k ] v) t
|
|
in
|
|
impl [] (StringMap.to_list connections)
|
|
|
|
let part2 connections =
|
|
let max_set = find_max_set connections in
|
|
List.sort compare max_set |> String.concat ","
|
|
|
|
let _ = Aoc.main load_file [ (Fun.id, part1); (Fun.id, part2) ]
|