red-black-tree.fs

[<Class>]
type 'a BinaryTree =
    member hd : 'a
    member left : 'a BinaryTree
    member right : 'a BinaryTree
    member exists : 'a -> bool
    member insert : 'a -> 'a BinaryTree
    member print : unit -> unit
    static member empty : 'a BinaryTree
    
type color = R | B    
type 'a tree =
    | E
    | T of color * 'a tree * 'a * 'a tree
    
module Tree =
    let hd = function
        | E -> failwith "empty"
        | T(c, l, x, r) -> x
        
    let left = function
        | E -> failwith "empty"
        | T(c, l, x, r) -> l
    
    let right = function
        | E -> failwith "empty"
        | T(c, l, x, r) -> r
        
    let rec exists item = function
        | E -> false
        | T(c, l, x, r) ->
            if item = x then true
            elif item < x then exists item l
            else exists item r
            
    let balance = function                              (* Red nodes in relation to black root *)
        | B, T(R, T(R, a, x, b), y, c), z, d            (* Left, left *)
        | B, T(R, a, x, T(R, b, y, c)), z, d            (* Left, right *)
        | B, a, x, T(R, T(R, b, y, c), z, d)            (* Right, left *)
        | B, a, x, T(R, b, y, T(R, c, z, d))            (* Right, right *)
            -> T(R, T(B, a, x, b), y, T(B, c, z, d))
        | c, l, x, r -> T(c, l, x, r)
    
    let insert item tree =
        let rec ins = function
            | E -> T(R, E, item, E)
            | T(c, a, y, b) as node ->
                if item = y then node
                elif item < y then balance(c, ins a, y, b)
                else balance(c, a, y, ins b)

        (* Forcing root node to be black *)                
        match ins tree with
            | E -> failwith "Should never return empty from an insert"
            | T(_, l, x, r) -> T(B, l, x, r)
            
    let rec print (spaces : int) = function
        | E -> ()
        | T(c, l, x, r) ->
            print (spaces + 4) r
            printfn "%s %A%A" (new System.String(' ', spaces)) c x
            print (spaces + 4) l
    
type 'a BinaryTree(inner : 'a tree) =
    member this.hd = Tree.hd inner
    member this.left = BinaryTree(Tree.left inner)
    member this.right = BinaryTree(Tree.right inner)
    member this.exists item = Tree.exists item inner
    member this.insert item = BinaryTree(Tree.insert item inner)
    member this.print() = Tree.print 0 inner
    static member empty = BinaryTree<'a>(E)