Difficulty in writing Red Black Tree in F#

Question

I am writing a red black tree in F#.

the code which I have written is below. I am facing 2 problems with this code

1. The rules of balancing the tree state that when the tree has a XYr or rXY type of imbalance I must recolor the 2 parent nodes and IF the grand parent node is not ROOT of the tree then it should be recolored as well.

The difficulty here is that in the recursive approach I only get the next node to work on.. so its hard to know what is the root node.

1. IN order to solve the above, I added another integer called height to my Node type (type node = Node of int * int * color). That made my pattern matching code in balanceTree function pretty long... but the problem is that when I recolor the grandparent the tree becomes imbalanced because the grand-grand-parent and grand-parent can be red in color which is not allowed.

Can someone recommend a clean way of resolving the issue.

type Color = 
  | R
  | B

type tree = 
  | Node of int * Color * tree * tree
  | Empty

let countNodes tree = 
    let rec incrCount = function
        | Empty -> 0
        | Node(_, _, n1, n2) -> 1 + (incrCount n1) + (incrCount n2)
    incrCount tree

let isTreeValid tree = 
  let getTreeBlackNodeHeight tree = 
    let rec getNodeHeight acc = function
      | Empty -> acc + 1
      | Node(_, R, n1, _) -> getNodeHeight acc n1
      | Node(_, B, n1, _) -> getNodeHeight (acc + 1) n1
    getNodeHeight 0 tree

  let isRootNodeBlack = function 
    | Empty -> true
    | Node(_, B, _, _) -> true
    | Node(_, R, _, _) -> false

  let rec areAllBlackHeightsSame height acc = function
    | Empty -> 
      if (acc + 1) = height then true else false
    | Node(_, R, n1, n2) -> areAllBlackHeightsSame height acc n1 && areAllBlackHeightsSame height acc n2
    | Node(_, B, n1, n2) ->  areAllBlackHeightsSame height (acc + 1) n1 && areAllBlackHeightsSame height (acc + 1) n2


  let allRedsMustHaveBlackChildren tree = 
    let getRootNodeColor = function
      | Empty -> Color.B
      | Node(_, y, _, _) -> y

    let rec checkChildColor = function
      | Empty -> true
      | Node(_, R, n1, n2) -> getRootNodeColor n1 = Color.B && getRootNodeColor n2 = Color.B && checkChildColor n1 && checkChildColor n2
      | Node(_, B, n1, n2) -> (checkChildColor n1) && (checkChildColor n2)
    checkChildColor tree

  (areAllBlackHeightsSame (getTreeBlackNodeHeight tree) 0 tree) && (isRootNodeBlack tree) && (allRedsMustHaveBlackChildren tree)

let insert x tree = 
    let rec createNode = function
        | Empty -> if (countNodes tree) = 0 then Node(x, B, Empty, Empty) else Node(x, R, Empty, Empty)
        | Node(i, c, n1, n2) when x > i -> Node(i, c, n1, (createNode n2)) 
        | Node(i, c, n1, n2) when x < i -> Node(i, c, (createNode n1), n2)
        | Node(i, _, _, _) when x = i -> failwith "Node already exists"
        | _ -> failwith "unknown"
    createNode tree

let colorToggle = function
    | (i, B) -> (i, R)
    | (i, R) -> (i, B)

let balanceTree tree =     
  let rec balance = function
      | Node(gpv, B, Node(p1v, R, Node(c1v, R, a, b), c), Node(p2v, R, d, e)) -> balance (Node(gpv, B, Node(p1v, B, Node(c1v, R, a, b), c), Node(p2v, B, d, e)))
      | Node(gpv, B, Node(p1v, R, a, Node(c2v, R, b, c)), Node(p2v, R, d, e)) -> balance (Node(gpv, B, Node(p1v, B, a, Node(c2v, R, b, c)), Node(p2v, B, e, e)))
      | Node(gpv, B, Node(p1v, R, a, b), Node(p2v, R, Node(c1v, R, c, d), e)) -> balance (Node(gpv, B, Node(p1v, B, a, b), Node(p2v, B, Node(c1v, R, c, d), e)))
      | Node(gpv, B, Node(p1v, R, a, b), Node(p2v, R, c, Node(c2v, R, d, e))) -> balance (Node(gpv, B, Node(p1v, B, a, b), Node(p2v, B, c, Node(c2v, R, d, e))))
      | Node(gpv, B, x4, Node(pv, R, x1, Node(cv, R, x2, x3))) -> balance (Node(pv, B, Node(gpv, R, x4, x1), Node(cv, R, x2, x3)))
      | Node(gpv, B, x4, Node(pv, R, Node(cv, R, x1, x2), x3)) -> balance (Node(pv, B, Node(gpv, R, x4, Node(cv, B, x1, x2)), x3))
      | Node(gpv, B, Node(pv, R, x1, Node(cv, R, x2, x3)), x4) -> balance (Node(pv, B, x1, Node(gpv, R, Node(cv, R, x2, x3), x4)))
      | Node(gpv, B, Node(pv, R, Node(cv, R, x1, x2), x3), x4) -> balance (Node(pv, B, (Node(cv, R, x1, x2)), Node(gpv, R, x3, x4)))
      | Node(i, x, n1, n2) -> Node(i, x, (balance n1), (balance n2))
      | Empty -> Empty
  balance tree

[<EntryPoint>]
let main args = 
  //let t1 = Node((35, B), Node((20, R), Node((10, B), Node((5, R), Empty, Empty), Empty), Node((25, B), Empty, Empty)), Node((85, R), Node((55, B), Node((40, R), Empty, Empty), Node((70, R), Empty, Empty)), Node((100, B), Empty, Empty)))
  let t2 = [1 .. 6] |> List.fold (fun acc i-> insert i acc) Empty
  printfn "Is Tree Valid : %b" (isTreeValid t2)
  let t3 = balanceTree t2
  printfn "is Tree Valid : %b" (isTreeValid t3)
  0

Answer 1

Standard ML-style implementation in F# looks like this:

type color = R | B
type 'a tree = E | T of color * 'a tree * 'a * 'a tree

let balance = function
  | B, T (R, T (R,a,x,b), y, c), z, d
  | B, T (R, a, x, T (R,b,y,c)), z, d
  | B, a, x, T (R, T (R,b,y,c), z, d)
  | B, a, x, T (R, b, y, T (R,c,z,d)) -> T (R, T (B,a,x,b), y, T (B,c,z,d))
  | col, a, x, b                      -> T (col, a, x, b) 

let insert x s = 
  let rec ins = function
    | E                  -> T (R,E,x,E)
    | T (col,a,y,b) as s ->
        if x < y then
          balance (col, ins a, y, b)
        elif x > y then
          balance (col, a, y, ins b)
        else
          s
  match ins s with
  | T (_,a,y,b) -> T (B,a,y,b)
  | t -> t