Files
linguist/samples/Standard ML/RedBlackTree.fun

255 lines
9.4 KiB
Standard ML

(* From Twelf *)
(* Red/Black Trees *)
(* Author: Frank Pfenning *)
functor RedBlackTree
(type key'
val compare : key' * key' -> order)
:> TABLE where type key = key' =
struct
type key = key'
type 'a entry = key * 'a
datatype 'a dict =
Empty (* considered black *)
| Red of 'a entry * 'a dict * 'a dict
| Black of 'a entry * 'a dict * 'a dict
type 'a Table = 'a dict ref
(* Representation Invariants *)
(*
1. The tree is ordered: for every node Red((key1,datum1), left, right) or
Black ((key1,datum1), left, right), every key in left is less than
key1 and every key in right is greater than key1.
2. The children of a red node are black (color invariant).
3. Every path from the root to a leaf has the same number of
black nodes, called the black height of the tree.
*)
local
fun lookup dict key =
let
fun lk (Empty) = NONE
| lk (Red tree) = lk' tree
| lk (Black tree) = lk' tree
and lk' ((key1, datum1), left, right) =
(case compare(key,key1)
of EQUAL => SOME(datum1)
| LESS => lk left
| GREATER => lk right)
in
lk dict
end
(* val restore_right : 'a dict -> 'a dict *)
(*
restore_right (Black(e,l,r)) >=> dict
where (1) Black(e,l,r) is ordered,
(2) Black(e,l,r) has black height n,
(3) color invariant may be violated at the root of r:
one of its children might be red.
and dict is a re-balanced red/black tree (satisfying all invariants)
and same black height n.
*)
fun restore_right (Black(e, Red lt, Red (rt as (_,Red _,_)))) =
Red(e, Black lt, Black rt) (* re-color *)
| restore_right (Black(e, Red lt, Red (rt as (_,_,Red _)))) =
Red(e, Black lt, Black rt) (* re-color *)
| restore_right (Black(e, l, Red(re, Red(rle, rll, rlr), rr))) =
(* l is black, deep rotate *)
Black(rle, Red(e, l, rll), Red(re, rlr, rr))
| restore_right (Black(e, l, Red(re, rl, rr as Red _))) =
(* l is black, shallow rotate *)
Black(re, Red(e, l, rl), rr)
| restore_right dict = dict
(* restore_left is like restore_right, except *)
(* the color invariant may be violated only at the root of left child *)
fun restore_left (Black(e, Red (lt as (_,Red _,_)), Red rt)) =
Red(e, Black lt, Black rt) (* re-color *)
| restore_left (Black(e, Red (lt as (_,_,Red _)), Red rt)) =
Red(e, Black lt, Black rt) (* re-color *)
| restore_left (Black(e, Red(le, ll as Red _, lr), r)) =
(* r is black, shallow rotate *)
Black(le, ll, Red(e, lr, r))
| restore_left (Black(e, Red(le, ll, Red(lre, lrl, lrr)), r)) =
(* r is black, deep rotate *)
Black(lre, Red(le, ll, lrl), Red(e, lrr, r))
| restore_left dict = dict
fun insert (dict, entry as (key,datum)) =
let
(* val ins : 'a dict -> 'a dict inserts entry *)
(* ins (Red _) may violate color invariant at root *)
(* ins (Black _) or ins (Empty) will be red/black tree *)
(* ins preserves black height *)
fun ins (Empty) = Red(entry, Empty, Empty)
| ins (Red(entry1 as (key1, datum1), left, right)) =
(case compare(key,key1)
of EQUAL => Red(entry, left, right)
| LESS => Red(entry1, ins left, right)
| GREATER => Red(entry1, left, ins right))
| ins (Black(entry1 as (key1, datum1), left, right)) =
(case compare(key,key1)
of EQUAL => Black(entry, left, right)
| LESS => restore_left (Black(entry1, ins left, right))
| GREATER => restore_right (Black(entry1, left, ins right)))
in
case ins dict
of Red (t as (_, Red _, _)) => Black t (* re-color *)
| Red (t as (_, _, Red _)) => Black t (* re-color *)
| dict => dict
end
(* function below from .../smlnj-lib/Util/int-redblack-set.sml *)
(* Need to check and improve some time *)
(* Sun Mar 13 08:22:53 2005 -fp *)
(* Remove an item. Returns true if old item found, false otherwise *)
local
exception NotFound
datatype 'a zipper
= TOP
| LEFTB of ('a entry * 'a dict * 'a zipper)
| LEFTR of ('a entry * 'a dict * 'a zipper)
| RIGHTB of ('a dict * 'a entry * 'a zipper)
| RIGHTR of ('a dict * 'a entry * 'a zipper)
in
fun delete t key =
let
fun zip (TOP, t) = t
| zip (LEFTB(x, b, z), a) = zip(z, Black(x, a, b))
| zip (LEFTR(x, b, z), a) = zip(z, Red(x, a, b))
| zip (RIGHTB(a, x, z), b) = zip(z, Black(x, a, b))
| zip (RIGHTR(a, x, z), b) = zip(z, Red(x, a, b))
(* bbZip propagates a black deficit up the tree until either the top
* is reached, or the deficit can be covered. It returns a boolean
* that is true if there is still a deficit and the zipped tree.
*)
fun bbZip (TOP, t) = (true, t)
| bbZip (LEFTB(x, Red(y, c, d), z), a) = (* case 1L *)
bbZip (LEFTR(x, c, LEFTB(y, d, z)), a)
| bbZip (LEFTB(x, Black(w, Red(y, c, d), e), z), a) = (* case 3L *)
bbZip (LEFTB(x, Black(y, c, Red(w, d, e)), z), a)
| bbZip (LEFTR(x, Black(w, Red(y, c, d), e), z), a) = (* case 3L *)
bbZip (LEFTR(x, Black(y, c, Red(w, d, e)), z), a)
| bbZip (LEFTB(x, Black(y, c, Red(w, d, e)), z), a) = (* case 4L *)
(false, zip (z, Black(y, Black(x, a, c), Black(w, d, e))))
| bbZip (LEFTR(x, Black(y, c, Red(w, d, e)), z), a) = (* case 4L *)
(false, zip (z, Red(y, Black(x, a, c), Black(w, d, e))))
| bbZip (LEFTR(x, Black(y, c, d), z), a) = (* case 2L *)
(false, zip (z, Black(x, a, Red(y, c, d))))
| bbZip (LEFTB(x, Black(y, c, d), z), a) = (* case 2L *)
bbZip (z, Black(x, a, Red(y, c, d)))
| bbZip (RIGHTB(Red(y, c, d), x, z), b) = (* case 1R *)
bbZip (RIGHTR(d, x, RIGHTB(c, y, z)), b)
| bbZip (RIGHTR(Red(y, c, d), x, z), b) = (* case 1R *)
bbZip (RIGHTR(d, x, RIGHTB(c, y, z)), b)
| bbZip (RIGHTB(Black(y, Red(w, c, d), e), x, z), b) = (* case 3R *)
bbZip (RIGHTB(Black(w, c, Red(y, d, e)), x, z), b)
| bbZip (RIGHTR(Black(y, Red(w, c, d), e), x, z), b) = (* case 3R *)
bbZip (RIGHTR(Black(w, c, Red(y, d, e)), x, z), b)
| bbZip (RIGHTB(Black(y, c, Red(w, d, e)), x, z), b) = (* case 4R *)
(false, zip (z, Black(y, c, Black(x, Red(w, d, e), b))))
| bbZip (RIGHTR(Black(y, c, Red(w, d, e)), x, z), b) = (* case 4R *)
(false, zip (z, Red(y, c, Black(w, Red(w, d, e), b))))
| bbZip (RIGHTR(Black(y, c, d), x, z), b) = (* case 2R *)
(false, zip (z, Black(x, Red(y, c, d), b)))
| bbZip (RIGHTB(Black(y, c, d), x, z), b) = (* case 2R *)
bbZip (z, Black(x, Red(y, c, d), b))
| bbZip (z, t) = (false, zip(z, t))
fun delMin (Red(y, Empty, b), z) = (y, (false, zip(z, b)))
| delMin (Black(y, Empty, b), z) = (y, bbZip(z, b))
| delMin (Black(y, a, b), z) = delMin(a, LEFTB(y, b, z))
| delMin (Red(y, a, b), z) = delMin(a, LEFTR(y, b, z))
| delMin (Empty, _) = raise Match
fun joinRed (Empty, Empty, z) = zip(z, Empty)
| joinRed (a, b, z) = let
val (x, (needB, b')) = delMin(b, TOP)
in
if needB
then #2(bbZip(z, Red(x, a, b')))
else zip(z, Red(x, a, b'))
end
fun joinBlack (a, Empty, z) = #2(bbZip(z, a))
| joinBlack (Empty, b, z) = #2(bbZip(z, b))
| joinBlack (a, b, z) = let
val (x, (needB, b')) = delMin(b, TOP)
in
if needB
then #2(bbZip(z, Black(x, a, b')))
else zip(z, Black(x, a, b'))
end
fun del (Empty, z) = raise NotFound
| del (Black(entry1 as (key1, datum1), a, b), z) =
(case compare(key,key1)
of EQUAL => joinBlack (a, b, z)
| LESS => del (a, LEFTB(entry1, b, z))
| GREATER => del (b, RIGHTB(a, entry1, z)))
| del (Red(entry1 as (key1, datum1), a, b), z) =
(case compare(key,key1)
of EQUAL => joinRed (a, b, z)
| LESS => del (a, LEFTR(entry1, b, z))
| GREATER => del (b, RIGHTR(a, entry1, z)))
in
(del(t, TOP); true) handle NotFound => false
end
end (* local *)
(* use non-imperative version? *)
fun insertShadow (dict, entry as (key,datum)) =
let val oldEntry = ref NONE (* : 'a entry option ref *)
fun ins (Empty) = Red(entry, Empty, Empty)
| ins (Red(entry1 as (key1, datum1), left, right)) =
(case compare(key,key1)
of EQUAL => (oldEntry := SOME(entry1);
Red(entry, left, right))
| LESS => Red(entry1, ins left, right)
| GREATER => Red(entry1, left, ins right))
| ins (Black(entry1 as (key1, datum1), left, right)) =
(case compare(key,key1)
of EQUAL => (oldEntry := SOME(entry1);
Black(entry, left, right))
| LESS => restore_left (Black(entry1, ins left, right))
| GREATER => restore_right (Black(entry1, left, ins right)))
in
(oldEntry := NONE;
((case ins dict
of Red (t as (_, Red _, _)) => Black t (* re-color *)
| Red (t as (_, _, Red _)) => Black t (* re-color *)
| dict => dict),
!oldEntry))
end
fun app f dict =
let fun ap (Empty) = ()
| ap (Red tree) = ap' tree
| ap (Black tree) = ap' tree
and ap' (entry1, left, right) =
(ap left; f entry1; ap right)
in
ap dict
end
in
fun new (n) = ref (Empty) (* ignore size hint *)
val insert = (fn table => fn entry => (table := insert (!table, entry)))
val insertShadow =
(fn table => fn entry =>
let
val (dict, oldEntry) = insertShadow (!table, entry)
in
(table := dict; oldEntry)
end)
val lookup = (fn table => fn key => lookup (!table) key)
val delete = (fn table => fn key => (delete (!table) key; ()))
val clear = (fn table => (table := Empty))
val app = (fn f => fn table => app f (!table))
end
end; (* functor RedBlackTree *)