diff options
Diffstat (limited to 'day7.ml')
-rw-r--r-- | day7.ml | 90 |
1 files changed, 90 insertions, 0 deletions
diff --git a/day7.ml b/day7.ml new file mode 100644 index 0000000..498d732 --- /dev/null +++ b/day7.ml @@ -0,0 +1,90 @@ +#use "topfind";; +#require "containers";; +#load "str.cma";; + +module StrMap = Map.Make(String) +module StrSet = Set.Make(String) +module IntSet = Set.Make(Int) +module Regex = Str + +let r = Regex.regexp {|\([a-z]+\) (\([0-9]+\))\($\| -> \(.+\)\)|} + +let parse_row s = + if Regex.string_match r s 0 then + let name = Regex.matched_group 1 s in + let weight = Regex.matched_group 2 s |> int_of_string in + let kids = + if Regex.matched_group 3 s = "" then [] + else CCString.split ~by:", " (Regex.matched_group 4 s) + in + (name, weight, kids) + else ("", 0, []) + +let input = List.map parse_row CCIO.(with_in "day7.txt" read_lines_l) + +let create_map f = + List.fold_left f StrMap.empty +let get_weights = + create_map (fun a (n, w, _) -> StrMap.add n w a) +let get_relations = + create_map (fun a (n, _, k) -> StrMap.add n k a) + +let weights = get_weights input +let relations = get_relations input + + +let rec total_node_weight node = + let own_weight = StrMap.find node weights in + let kids = StrMap.find node relations in + let kids_weights = List.map total_node_weight kids in + own_weight + List.fold_left (+) 0 kids_weights + + +let find_different_weights kids = + if CCList.is_empty kids then None + else + let kids_weights = List.map total_node_weight kids in + let w_set = IntSet.of_list kids_weights in + if IntSet.cardinal w_set = 1 then None + else + let min_w = IntSet.min_elt w_set in + let max_w = IntSet.max_elt w_set in + let count v = + kids_weights + |> List.filter (fun w -> w = v) + |> List.length in + let min_count = count min_w in + let max_count = count max_w in + let (wrong_w, right_w) = + if min_count < max_count then (min_w, max_w) else (max_w, min_w) in + let wrong_i, _ = + CCList.find_idx (fun x -> x = wrong_w) kids_weights + |> Option.value ~default:(0, 0) in + let wrong_kid = List.nth kids wrong_i in + Some (wrong_kid, right_w) + + +let rec find_correct_weight node right_w = + let kids = StrMap.find node relations in + match find_different_weights kids with + | Some (wrong_kid, right_w) -> + find_correct_weight wrong_kid right_w + | None -> + let kids_weights = List.map total_node_weight kids in + let kids_total_w = List.fold_left (+) 0 kids_weights in + right_w - kids_total_w + + +let create_set f = + List.fold_left f StrSet.empty +let get_nodes = + create_set (fun a (n, _, _) -> StrSet.add n a) +let get_kids = + create_set (fun a (_, _, k) -> StrSet.union a (StrSet.of_list k)) + +let nodes = get_nodes input +let kids = get_kids input +let root = StrSet.diff nodes kids |> StrSet.choose + +let part_1 = root |> Printf.printf "%s\n" +let part_2 = find_correct_weight root 0 |> Printf.printf "%d\n" |