Red-Black Tree

red_black_tree.sml

(* MLton用 ORD_KEY signature *)
signature ORD_KEY =
sig
  type ord_key
  val compare : (ord_key * ord_key) -> order;
end;

(* 赤黒木のシグネチャ *)
signature RED_BLACK_TREE =
sig
  exception EmptyTree;
  
  type item;
  datatype tree =     Empty 
                | BlackNode of item * tree * tree 
                |   RedNode of item * tree * tree;
  
  val empty   :  tree;
  val isEmpty :  tree -> bool;
  val member  : (tree * item) -> bool;
  val height  :  tree -> int;
  val insert  : (tree * item) -> tree;
  val delete  : (tree * item) -> tree;
  
  val app       : (item -> unit) -> tree -> unit;
  val exists    : (item -> bool) -> tree -> bool;
  val numItems  : tree -> int;
  val listItems : tree -> item list;
end;

(* 赤黒木の実装 *)
functor RedBlackTreeFunctor
  (
    OrderedKey : ORD_KEY
  ) :> RED_BLACK_TREE
  where type item = OrderedKey.ord_key
 =
struct
  exception EmptyTree;
  
  type item = OrderedKey.ord_key;
  datatype tree =     Empty 
                | BlackNode of item * tree * tree 
                |   RedNode of item * tree * tree;
  
  val empty = Empty;
  
  fun isEmpty Empty = true
    | isEmpty _     = false;
  
  fun member (Empty, _) = false
    | member (BlackNode(value, left, right), query) = member_aux ((value, left, right), query)
    | member (  RedNode(value, left, right), query) = member_aux ((value, left, right), query)
  and member_aux ((value, left, right), query) = 
    case (OrderedKey.compare(query, value))
      of EQUAL   => true
       | LESS    => member (left,  query)
       | GREATER => member (right, query);
  
  fun height Empty = 0
    | height (BlackNode(_, left, right)) = height_aux (left, right)
    | height (  RedNode(_, left, right)) = height_aux (left, right)
  and height_aux (left, right) = 
    let
      fun max (x, y) = if x > y then x else y;
    in 
      1 + max(height left, height right)
    end;
  
  fun balance Empty = Empty
    | balance (BlackNode(u, RedNode(v, RedNode(w, t1, t2), t3), t4)) = RedNode(v, BlackNode(w, t1, t2), BlackNode(u, t3, t4))
    | balance (BlackNode(u, RedNode(v, t1, RedNode(w, t2, t3)), t4)) = RedNode(w, BlackNode(v, t1, t2), BlackNode(u, t3, t4))
    | balance (BlackNode(u, t1, RedNode(v, RedNode(w, t2, t3), t4))) = RedNode(w, BlackNode(u, t1, t2), BlackNode(v, t3, t4))
    | balance (BlackNode(u, t1, RedNode(v, t2, RedNode(w, t3, t4)))) = RedNode(v, BlackNode(u, t1, t2), BlackNode(w, t3, t4))
    | balance node = node;
  
  fun insert (tree, query) = let
    fun to_black (BlackNode(value, left, right)) = BlackNode (value, left, right)
      | to_black (  RedNode(value, left, right)) = BlackNode (value, left, right);
  in to_black (insert_aux (tree, query)) end
  and insert_aux (Empty, query) = RedNode(query, Empty, Empty)
    | insert_aux (BlackNode(value, left, right), query) =
      let
        val node = 
          case (OrderedKey.compare(query, value))
            of EQUAL   => BlackNode(value, left, right)
             | LESS    => BlackNode(value, insert_aux (left, query), right)
             | GREATER => BlackNode(value, left, insert_aux (right, query));
      in balance node end
    | insert_aux (  RedNode(value, left, right), query) = 
      let
        val node = 
          case (OrderedKey.compare(query, value))
            of EQUAL   =>   RedNode(value, left, right)
             | LESS    =>   RedNode(value, insert_aux (left, query), right)
             | GREATER =>   RedNode(value, left, insert_aux (right, query));
      in node end;
  
  fun delete (tree, query) = let
    fun to_black (BlackNode(value, left, right)) = BlackNode (value, left, right)
      | to_black (  RedNode(value, left, right)) = BlackNode (value, left, right);
  in to_black (delete_aux (tree, query)) end
  and delete_aux (Empty, query) = Empty
    | delete_aux (BlackNode(value, left, right), query) = let
        val node = 
          case (OrderedKey.compare(query, value))
            of EQUAL   => (case (right) of Empty => Empty | _ => BlackNode(search_min right, left, delete_min right))
             | LESS    => BlackNode(value, delete_aux (left, query), right)
             | GREATER => BlackNode(value, left, delete_aux (right, query));
    in balance node end
    | delete_aux (  RedNode(value, left, right), query) = let
        val node = 
          case (OrderedKey.compare(query, value))
            of EQUAL   => (case (right) of Empty => Empty | _ =>   RedNode(search_min right, left, delete_min right))
             | LESS    =>   RedNode(value, delete_aux (left, query), right)
             | GREATER =>   RedNode(value, left, delete_aux (right, query));
    in node end
  and search_min Empty = raise EmptyTree
    | search_min (BlackNode(value, Empty, right)) = value
    | search_min (  RedNode(value, Empty, right)) = value
    | search_min (BlackNode(value,  left, right)) = search_min left
    | search_min (  RedNode(value,  left, right)) = search_min left
  and delete_min Empty = Empty
    | delete_min (BlackNode(value, Empty, right)) = right
    | delete_min (  RedNode(value, Empty, right)) = right
    | delete_min (BlackNode(value,  left, right)) = BlackNode(value, delete_min left, right)
    | delete_min (  RedNode(value,  left, right)) =   RedNode(value, delete_min left, right)
  (*
  and delete_aux (Empty, query) = Empty
    | delete_aux (BlackNode(value, left, right), query) = let
        val node = 
          case (OrderedKey.compare(query, value))
            of EQUAL   => (case (left) of Empty => Empty | _ => BlackNode(search_max left, delete_max left, right))
             | LESS    => BlackNode(value, delete_aux (left, query), right)
             | GREATER => BlackNode(value, left, delete_aux (right, query));
    in balance node end
    | delete_aux (  RedNode(value, left, right), query) = let
        val node = 
          case (OrderedKey.compare(query, value))
            of EQUAL   => (case (left) of Empty => Empty | _ =>   RedNode(search_max left, delete_max left, right))
             | LESS    =>   RedNode(value, delete_aux (left, query), right)
             | GREATER =>   RedNode(value, left, delete_aux (right, query));
    in node end
  and search_max Empty = raise EmptyTree
    | search_max (BlackNode(value, left, Empty)) = value
    | search_max (  RedNode(value, left, Empty)) = value
    | search_max (BlackNode(value, left, right)) = search_max right
    | search_max (  RedNode(value, left, right)) = search_max right
  and delete_max Empty = Empty
    | delete_max (BlackNode(value, left, Empty)) = left
    | delete_max (  RedNode(value, left, Empty)) = left
    | delete_max (BlackNode(value, left, right)) = BlackNode(value, left, delete_max right)
    | delete_max (  RedNode(value, left, right)) =   RedNode(value, left, delete_max right)
  *);
  
  fun app pred Empty  = ()
    | app pred (BlackNode(value, left, right)) = app_aux pred (value, left, right) 
    | app pred (  RedNode(value, left, right)) = app_aux pred (value, left, right)
  and app_aux pred (value, left, right) = (app pred left; pred value; app pred right);
  
  fun exists pred Empty  = false
    | exists pred (BlackNode(value, left, right)) = exists_aux pred (value, left, right)
    | exists pred (  RedNode(value, left, right)) = exists_aux pred (value, left, right)
  and exists_aux pred (value, left, right) = if (pred value) then true else ((exists pred left) orelse (exists pred right));
  
  fun numItems Empty = 0
    | numItems (BlackNode(value, left, right)) = 1 + (numItems left) + (numItems right)
    | numItems (  RedNode(value, left, right)) = 1 + (numItems left) + (numItems right);
  
  fun listItems Empty = nil
    | listItems (BlackNode(value, left, right)) = (listItems left) @ [ value ] @ (listItems right)
    | listItems (  RedNode(value, left, right)) = (listItems left) @ [ value ] @ (listItems right);
end;

(* 比較付き整数 *)
structure OrderedInt :> ORD_KEY 
  where type ord_key = Int.int
= struct
  type ord_key = Int.int;
  
  fun compare (lhs, rhs) = 
    if lhs = rhs
      then EQUAL
      else 
        if lhs < rhs then LESS else GREATER; 
end;

(* 赤黒木(整数) *)
structure IntRedBlackTree = RedBlackTreeFunctor
  (
    OrderedInt
  );

(* デバッグ用ユーティリティ *)
fun to_string IntRedBlackTree.Empty = "_"
  | to_string (IntRedBlackTree.BlackNode(value, left, right)) = to_string_aux ("Blk", value, left, right)
  | to_string (IntRedBlackTree.RedNode  (value, left, right)) = to_string_aux ("Red", value, left, right)
and to_string_aux (color, value, left, right) = "(" ^ (String.concatWith ", " [color, Int.toString value, to_string left, to_string right]) ^ ")";

fun tree_to_dot tree = let
  fun to_string  IntRedBlackTree.Empty                   = Option.NONE
    | to_string (IntRedBlackTree.BlackNode(value, _, _)) = Option.SOME (Int.toString value)
    | to_string (IntRedBlackTree.RedNode  (value, _, _)) = Option.SOME (Int.toString value);
  
  fun build_black_node value = "  " ^ (Int.toString value) ^ " [fontcolor=white,fillcolor=black,style=filled];\n";
  fun build_red_node   value = "  " ^ (Int.toString value) ^ " [fontcolor=white,fillcolor=red,style=filled];\n";
  
  fun build_arrow (src, dst) = let
    fun arrow (src, dst) = "  " ^ src ^ " -> " ^ dst ^ ";\n";
  in
    case (dst)
      of (Option.SOME dst) => arrow (src, dst)
       |  Option.NONE      => ""
  end;
  
  fun dot  IntRedBlackTree.Empty = ""
    | dot (IntRedBlackTree.BlackNode(src, left, right)) = (build_black_node src) ^ (dot_aux (src, left, right))
    | dot (IntRedBlackTree.RedNode  (src, left, right)) = (build_red_node   src) ^ (dot_aux (src, left, right))
  and dot_aux (src, left, right) = let
      val lhs = build_arrow (Int.toString src, to_string left);
      val rhs = build_arrow (Int.toString src, to_string right);
    in (lhs ^ rhs) ^ (dot left) ^ (dot right) end;
in "digraph test {\n" ^ (dot tree) ^ "}" end;

(* 実験 *)
(*
val tree = IntRedBlackTree.empty;

val tree = IntRedBlackTree.insert (tree,  1);
val tree = IntRedBlackTree.insert (tree,  2);
val tree = IntRedBlackTree.insert (tree,  3);
val tree = IntRedBlackTree.insert (tree,  4);
val tree = IntRedBlackTree.insert (tree,  5);
val tree = IntRedBlackTree.insert (tree,  6);
val tree = IntRedBlackTree.insert (tree,  7);
val tree = IntRedBlackTree.insert (tree,  8);
val tree = IntRedBlackTree.insert (tree,  9);
val tree = IntRedBlackTree.insert (tree, 10);
val tree = IntRedBlackTree.insert (tree, 11);
val tree = IntRedBlackTree.insert (tree, 12);

val () = print ((tree_to_dot tree) ^ "\n");
*)