Arn.ml

Ce module présente une implantation classique des arbres rouges/noires qui suit la représentation du cours d'algo. La seule particularité est l'utilisation d'une fonction (de première classe) permettant d'enregistrer les appels aux transformations d'arbres (rotation.)

(* Arbre rouge/noir *)

(* Les clefs sont des entiers *)
type t =
    Empty
  | Node of int * bool * t * t

(* Arbres vides *)
let empty = Empty
let is_empty t = t = Empty

(* Le noeud est-il rouge ? *)
let is_red = function
    Node(_,true,_,_) -> true
  | _ -> false

(* Recherche d'une clef *)
let rec mem k = function
    Empty -> false
  | Node(k',_,_,_) when k=k' -> true
  | Node(k',_,fg,_) when k<k'-> mem k fg
  | Node(_,_,_,fd) -> mem k fd

(* Affichage classique sur la sortie standard *)
let rec print = function
    Empty -> Format.printf "()"
  | Node(k,c,fg,fd) ->
      begin
	Format.printf "@,(@[<v>%i,%b,@," k c;
	print fg;
	Format.printf ",";
	print fd;
	Format.printf "@]@,)"
      end

(* Affichage avec tikz *)
let rec print_tikz = function
    Empty -> Format.printf "child {node {\ }}@,"
  | Node(k,true,fg,fd) ->
      begin
	Format.printf "child {@[<v>node[red_nd] {%i}@," k;
	print_tikz fg;
	print_tikz fd;
	Format.printf "@]}@,";
      end
  | Node(k,_,fg,fd) ->
      begin
	Format.printf "child {@[<v>node {%i}@," k;
	print_tikz fg;
	print_tikz fd;
	Format.printf "@]}@,";
      end

(* Affichage avec pour LaTeX (a inserer dans un fichier .tex) *)
let print_tex t =
  begin
    Format.printf "\\tikzstyle{red_nd}=[draw=red,circle]@.";
    Format.printf "\\begin{center}\\begin{tikzpicture}";
    Format.printf "[every node/.style={draw=black,circle}]@.";
    begin
      match t with
	  Empty -> Format.printf "\\node {\ };@."
	| Node(k,_,fg,fd) ->
	    begin
	      Format.printf "\\node@[<v> {%i}
[level/.style={level distance={max(8mm,30mm-(5mm*#1))},sibling distance=260mm/(2^#1)}] @," k;
	      print_tikz fg;
	      print_tikz fd;
	      Format.printf "@];@.";
	    end
    end;
    Format.printf "\\end{tikzpicture}\\end{center}@."
  end

(* Transformations *)

(* eclatement *)
let eclate update = function
    Node(k,_,Node(kg,_,fgg,fdg),Node(kd,_,fgd,fdd)) ->
      begin
	ignore(update "eclate");
	Node(k,true,Node(kg,false,fgg,fdg),Node(kd,false,fgd,fdd))
      end
  | _ -> assert false

(* rotation droite *)
let rd update = function
    Node(k,_,Node(kg,_,fgg,fgd),fd) ->
      begin
	ignore(update "rd");
	Node(kg,false,fgg,Node(k,true,fgd,fd))
      end
  | _ -> assert false

(* rotation gauche *)
let rg update = function
    Node(k,_,fg,Node(kd,_,fdg,fdd)) ->
      begin
	ignore(update "rg");
	Node(kd,false,Node(k,true,fg,fdg),fdd)
      end
  | _ -> assert false

(* rotation gauche-droite *)
let rgd update = function
    Node(k,_,Node(kg,c,fgg,Node(kgd,_,fgdg,fgdd)),fd) ->
      begin
	ignore(update "rgd");
	Node(kgd,false,Node(kg,c,fgg,fgdd),Node(k,true,fgdd,fd))
      end
  | _ -> assert false

(* rotation droite-gauche *)
let rdg update = function
    Node(k,_,fg,Node(kd,c,Node(kdg,_,fdgg,fdgd),fdd)) ->
      begin
	ignore(update "rdg");
	Node(kdg,false,Node(k,true,fg,fdgg),Node(kd,c,fdgd,fdd))
      end
  | _ -> assert false

(* insertion dans un arn (fonction recursive) *)
let rec insert update k = function
    Empty -> (1,Node(k,true,Empty,Empty))
  | Node(k',color,fg,fd) as n when k'=k -> (0,n)
  | Node(k',color,fg,fd) when k<k' ->
      begin
	let (r,nt) = insert update k fg in
	let n = Node(k',color,nt,fd) in
	  if color then
	    (1+r,n)
	  else
	    begin
	      if abs r = 2 then
		if is_red fd then
		  (1,eclate update n)
		else
		  let nnt =
		    if (r = 2) then rd update n
		    else rgd update n
		  in
		    (0,nnt)
	      else
		(0,Node(k',color,nt,fd))
	    end
      end
  | Node(k',color,fg,fd) ->
      begin
	let (r,nt) = insert update k fd in
	let n = Node(k',color,fg,nt) in
	  if color then
	    (1-3*r,n)
	  else
	    begin
	      if abs r = 2 then
		if is_red fg then
		  (1,eclate update n)
		else
		  let nnt =
		    if (r = -2) then rg update n
		    else rdg update n
		  in
		    (0,nnt)
	      else
		(0,Node(k',color,fg,nt))
	    end
      end

(* insertion dans un ARN (fonction d'appel) *)
let add update k t =
  match insert update k t with
      (1,Node(k',_,fg,fd)) -> Node(k',false,fg,fd)
    | (_,nt) -> nt

Ce module est similaire au module avl.ml, les même remarques s'y appliquent.