113 lines
3.6 KiB
OCaml
113 lines
3.6 KiB
OCaml
(* Red-black trees according to the following classic paper:
|
|
|
|
Chris Okasaki, Red-Black Trees in a Functional
|
|
Setting. J. Funct. Program. 9(4): 471-477 (1999)
|
|
*)
|
|
|
|
type colour = Red | Black
|
|
|
|
type 'a t =
|
|
Ext
|
|
| Int of colour * 'a t * 'a * 'a t
|
|
|
|
let empty = Ext
|
|
|
|
let is_empty m = (m = empty)
|
|
|
|
let blacken = function
|
|
Ext -> Ext
|
|
| Int (_, left, root, right) -> Int (Black, left, root, right)
|
|
|
|
let balance colour left root right =
|
|
match colour, left, root, right with
|
|
Black, Int (Red, Int (Red, a, x, b), y, c), z, d
|
|
| Black, Int (Red, a, x, Int (Red, b, y, c)), z, d
|
|
| Black, a, x, Int (Red, Int (Red, b, y, c), z, d)
|
|
| Black, a, x, Int (Red, b, y, Int (Red, c, z, d)) ->
|
|
Int (Red, Int (Black, a, x, b), y, Int (Black, c, z, d))
|
|
| _ -> Int (colour, left, root, right)
|
|
|
|
type choice = Old | New
|
|
|
|
let choose ~old ~new' = function
|
|
Old -> old
|
|
| New -> new'
|
|
|
|
exception Physical_equality
|
|
|
|
let add ~cmp choice elt tree =
|
|
let rec insert = function
|
|
Ext -> Int (Red, Ext, elt, Ext) (* A leaf *)
|
|
| Int (colour, left, root, right) ->
|
|
let diff = cmp elt root in
|
|
if diff = 0 then
|
|
let root' = choose ~new':elt ~old:root choice
|
|
in if root == root' then raise Physical_equality
|
|
else Int (colour, left, root', right)
|
|
else if diff < 0 then
|
|
balance colour (insert left) root right
|
|
else balance colour left root (insert right)
|
|
in try blacken (insert tree) with
|
|
Physical_equality -> tree
|
|
|
|
let remove : type a b . cmp:(a -> b -> int) -> a -> b t -> b t = fun ~cmp elt tree ->
|
|
(* TODO: this leaves the tree not properly balanced. *)
|
|
let rec bst_shift_up : b t -> b t = function
|
|
| Ext -> failwith "unknown error"
|
|
| Int (colour, left, root, right) ->
|
|
(
|
|
ignore root; (* we delete the root *)
|
|
match left, right with
|
|
| Ext, Ext -> Ext
|
|
| Ext, Int (_rcolour, _rleft, rroot, _rright) ->
|
|
let new_right = bst_shift_up right in
|
|
Int (colour, Ext, rroot, new_right)
|
|
| Int (_lcolour, _lleft, lroot, _lright), _ ->
|
|
let new_left = bst_shift_up left in
|
|
Int (colour, new_left, lroot, right)
|
|
) in
|
|
let rec bst_delete : a -> b t -> b t = fun elt -> function
|
|
| Ext -> failwith "remove in red-black tree: element not found"
|
|
| Int (colour, left, root, right) as current ->
|
|
let c = cmp elt root in
|
|
if c = 0 then bst_shift_up current
|
|
else if c < 0 then Int (colour, bst_delete elt left, root, right)
|
|
else Int (colour, left, root, bst_delete elt right)
|
|
in
|
|
bst_delete elt tree
|
|
|
|
exception Not_found
|
|
|
|
let rec find ~cmp elt = function
|
|
Ext -> raise Not_found
|
|
| Int (_, left, root, right) ->
|
|
let diff = cmp elt root in
|
|
if diff = 0 then root
|
|
else if diff < 0 then find ~cmp elt left
|
|
else find ~cmp elt right
|
|
|
|
let find_opt ~cmp elt tree =
|
|
try Some (find ~cmp elt tree) with Not_found -> None
|
|
|
|
(* Inorder iterators *)
|
|
|
|
let rec iter f = function
|
|
Ext -> ()
|
|
| Int (_, left, root, right) -> iter f left; f root; iter f right
|
|
|
|
let rec inorder acc = function
|
|
Ext -> acc
|
|
| Int (_, left, root, right) -> inorder (root :: inorder acc right) left
|
|
|
|
let elements t = inorder [] t
|
|
|
|
let rec fold_inc f ~init = function
|
|
Ext -> init
|
|
| Int (_, left, root, right) ->
|
|
fold_inc f ~init:(f ~elt:root ~acc:(fold_inc f ~init left)) right
|
|
|
|
let rec fold_dec f ~init = function
|
|
Ext -> init
|
|
| Int (_, left, root, right) ->
|
|
fold_dec f ~init:(f ~elt:root ~acc:(fold_dec f ~init right)) left
|