1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
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"
|