update lecture 08 notes

lectures/08/rbtree.ml

@@ -1,23 +1,47 @@
 type color = R | B
-
-type 'a rbtree = L | N of color * 'a rbtree * 'a * 'a rbtree
+type 'a t = E | N of (color * 'a t * 'a * 'a t)
 
 let balance = function
   | B, N (R, N (R, a, x, b), y, c), z, d
   | B, N (R, a, x, N (R, b, y, c)), z, d
   | B, a, x, N (R, N (R, b, y, c), z, d)
   | B, a, x, N (R, b, y, N (R, c, z, d)) ->
-    N (R, N (B, a, x, b), y, N (B, c, z, d))
+      N (R, N (B, a, x, b), y, N (B, c, z, d))
   | c, l, x, r ->
-    N (c, l, x, r)
+      N (c, l, x, r)
 
-let insert x t =
+let rec insert x t =
   let rec ins = function
-    | L -> N (R, L, x, L)
+    | E -> N (R, E, x, E)
     | N (c, l, y, r) when x < y ->
-      balance (c, ins l, x, r)
+        balance (c, ins l, y, r)
     | N (c, l, y, r) when x > y ->
-      balance (c, l, x, ins r)
-    | _ -> t
+        balance (c, l, y, ins r)
+    | t -> t
   in
-  ins t
+  match ins t with
+  | E -> failwith "insert: impossible"
+  | N (_, l, y, r) -> N (B, l, y, r)
+
+let of_list l =
+  List.fold_left (Fun.flip insert) E l
+
+let print pr t =
+  let rec aux level t =
+    Printf.printf "%*s" (3 * level) "";
+    match t with
+    | E -> Printf.printf "-\n"
+    | N (c, l, x, r) ->
+        if c = R then (
+          Printf.printf "\u{1b}[1;31m";
+          pr x;
+          Printf.printf "\u{1b}[1;0m\n"
+        )
+        else (
+          pr x;
+          Printf.printf "\n"
+        );
+        aux (level + 1) l;
+        aux (level + 1) r
+  in
+  aux 0 t