algorithm - F#写红黑树的难点

Question

我正在用 F# 编写一个红黑树。

我写的代码如下。我在使用这段代码时遇到了 2 个问题

1. 平衡树的规则规定，当树有 XYr 或 rXY 类型的不平衡时，我必须重新着色 2 个父节点，如果祖父节点不是树的 ROOT，那么它也应该重新着色。

这里的困难在于，在递归方法中，我只能获取下一个节点来处理……所以很难知道什么是根节点。

1. 为了解决上述问题，我在我的 Node 类型中添加了另一个称为 height 的整数(type node = Node of int * int * color)。这使得我在 balanceTree 函数中的模式匹配代码很长......但问题是，当我重新给祖 parent 着色时，树变得不平衡，因为祖 parent 和祖 parent 可以是红色的，这是不允许的。

有人可以推荐一种干净的方法来解决这个问题。

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

F# 中的标准 ML 风格实现如下所示:

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