summary refs log tree commit diff stats
path: root/day7.ml
diff options
context:
space:
mode:
Diffstat (limited to 'day7.ml')
-rw-r--r--day7.ml90
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"