day 5 part 2

5/main.ml

@@ -10,19 +10,16 @@ let rec list_of_lines in_file =
     close_in in_file;
     []
 
-let is_not_empty lst = 
-  match lst with
-  | [] -> false 
-  | _  -> true
-
 let rec parse_int_list sep str =
   let list1 = (List.filter (fun a -> length a != 0) (String.split_on_char sep str)) in
   List.map int_of_string list1
 
+(* Ordering map is a map that indexes int and maps it into a set of ints,
+  the indexed element must be before each of the elements of the mapped set*)
 module IntMap = Map.Make(Int)
 module IntSet = Set.Make(Int)
 
-let rec parse_ordering lines must_be_before =
+let rec parse_ordering lines ordering_map =
   match lines with
   | head :: tail ->
     if length head > 0 then
@@ -33,45 +30,61 @@ let rec parse_ordering lines must_be_before =
         match prev with 
         | Some set -> Some (IntSet.add e1 set)
         | None -> Some (IntSet.add e1 IntSet.empty) in
-      let new_map = IntMap.update e2 update_fun must_be_before in
+      let new_map = IntMap.update e2 update_fun ordering_map in
       parse_ordering tail new_map
     else 
-      (tail, must_be_before)
+      (tail, ordering_map)
-  | _ -> (lines, must_be_before)
+  | _ -> (lines, ordering_map)
 
 let parse in_file =
   let lines = list_of_lines in_file in
-  let lines_rest, must_be_before = parse_ordering lines IntMap.empty in
+  let lines_rest, ordering_map = parse_ordering lines IntMap.empty in
   let entries = List.map (fun str -> List.rev @@ parse_int_list ',' str) lines_rest in
-  (must_be_before, entries)
+  (ordering_map, entries)
 
-let rec is_good must_be_before already_appeared entry =
+let rec is_good ordering_map already_appeared entry =
   match entry with
   | head :: tail -> begin
     let new_already_appeared = IntSet.add head already_appeared in
-    match IntMap.find_opt head must_be_before with
+    match IntMap.find_opt head ordering_map with
     | Some set ->
       let intersect = IntSet.inter set new_already_appeared in
       if IntSet.is_empty intersect then
-        is_good must_be_before new_already_appeared tail
+        is_good ordering_map new_already_appeared tail
       else
         false
     | None ->
-      is_good must_be_before new_already_appeared tail
+      is_good ordering_map new_already_appeared tail
     end
   | _ -> true
 
-let rec summarise_good_entries must_be_before entries =
+let get_middle_element lst =
-  match entries with
+  List.nth lst (List.length lst / 2)
-  | head :: tail ->
+
-    if is_good must_be_before IntSet.empty head then
+let ordering_less_than ordering_map x1 x2 =
-      List.nth head (List.length head / 2) + summarise_good_entries must_be_before tail
+  match IntMap.find_opt x1 ordering_map with
+  | Some x1_set -> x1_set |> IntSet.find_opt x2 |> Option.is_some
+  | None -> false
+
+let ordering_function ordering_map x1 x2 =
+  if ordering_less_than ordering_map x1 x2 then (-1)
+  else if ordering_less_than ordering_map x2 x1 then 1
+  else 0
+
+let fix_ordering ordering_map lst =
+  List.sort (ordering_function ordering_map) lst
+
+let rec summarise_entries ordering_map acc elem =
+  if is_good ordering_map IntSet.empty elem then
+    let v = get_middle_element elem in
+    (fst acc + v, snd acc)
   else
-      summarise_good_entries must_be_before tail
+    let new_elem = fix_ordering ordering_map elem in
-  | _ -> 0
+    let v = get_middle_element new_elem in
+    (fst acc, snd acc + v)
 
 let () =
   let f = open_in "input.txt" in
-  let must_be_before, entries = parse f in
+  let ordering_map, entries = parse f in
-  let result = summarise_good_entries must_be_before entries in
+  let result, result2 = List.fold_left (summarise_entries ordering_map) (0, 0) entries in
-  printf "%d\n" result
+  printf "%d\n%d\b" result result2