golang-redBlackTree-insert

package main

import (
    "fmt"
)

func main() {
    fmt.Println("red-black-tree")
    var data = []int{6, 5, 3, 1, 8, 7, 2, 4, 9, 0, 3}
    //var data = []int{6, 5, 3, 1}

    tree := NewTree()
    for _, v := range data {
        tree.insertB(v)
    }
    fmt.Println(tree.size)
}

type node struct {
    l, r, p *node
    v       int
    c       bool //0 = red 1=black
}

type Tree struct {
    root *node
    last *node
    size int
}

func NewTree() *Tree {
    return &Tree{}
}

func (n *node) getGp() *node {
    if n.p == nil {
        return nil
    }
    if n.p.p == nil {
        return nil
    }
    return n.p.p
}

func (n *node) getUn() *node {
    if n.getGp() == nil {
        return nil
    }
    if n.p == n.getGp().l {
        return n.getGp().r
    } else {
        return n.getGp().l
    }
}

// rotate cur node
func (tree *Tree) rotateR(n *node) {
    gp := n.getGp()
    p := n.p
    r := n.r

    n.r = p
    n.p = gp

    if p != nil {
        p.p = n
        p.l = r
        if r != nil {
            r.p = p
        }
    }
    if tree.root == p {
        tree.root = n
    }

    if gp != nil {
        if gp.l == p {
            gp.l = n
        } else {
            gp.r = n
        }
    }
}

func (tree *Tree) rotateL(n *node) {
    gp := n.getGp()
    p := n.p
    l := n.l

    n.l = p
    n.p = gp

    if p != nil {
        p.p = n
        p.r = l
        if l != nil {
            l.p = p
        }
    }

    if tree.root == p {
        tree.root = n
    }

    if gp != nil {
        if gp.l == p {
            gp.l = n
        } else {
            gp.r = n
        }
    }

}

func getMin() {

}

func getMax() {

}

func (tree *Tree) insertB(v int) {
    if tree.root == nil {
        tree.root = new(node)
        tree.root.v = v
        tree.root.c = true
        tree.size++
    }

    if v < tree.root.v {
        if tree.Insert(&tree.root.l, v, tree.root) {
            printT(tree.root)
            tree.size++
            tree.inserCase(tree.last)
        }
    }

    if v > tree.root.v {
        if tree.Insert(&tree.root.r, v, tree.root) {
            printT(tree.root)
            tree.size++
            tree.inserCase(tree.last)
        }
    }
    printT(tree.root)
}

func (tree *Tree) Insert(n **node, v int, fa *node) bool {

    pn := (*n)
    if (*n) == nil {
        (*n) = new(node)
        (*n).v = v
        (*n).p = fa
        tree.last = (*n)
        return true
    }

    if v > pn.v {
        tree.Insert(&(pn.r), v, *n)
    }

    if v < pn.v {
        tree.Insert(&(pn.l), v, *n)
    }

    if v == pn.v {
        return false
    }
    return true
}

func (tree *Tree) inserCase(n *node) {
    if n.p == nil {
        n.c = true
        tree.root = n
        return
    }

    if n.p.c == false {
        if n.getUn() != nil && n.getUn().c == false {
            fmt.Println("aaaaa", n)
            n.p.c = true
            n.getUn().c = true
            n.getGp().c = false
            tree.inserCase(n.getGp()) //if root node
        } else {
            //nil or black
            if n == n.p.r && n.getGp() != nil && n.p == n.getGp().l {
                tree.rotateL(n)
                tree.rotateR(n)
                n.c = true
                n.l.c = false
                n.r.c = false
            }
            if n == n.p.l && n.getGp() != nil && n.p == n.getGp().r {
                tree.rotateR(n)
                tree.rotateL(n)
                n.c = true
                n.l.c = false
                n.r.c = false
            }
            if n == n.p.l && n.getGp() != nil && n.p == n.getGp().l {

                n.p.c = true
                if n.getGp() != nil {
                    n.getGp().c = false
                }
                tree.rotateR(n.p)
            }
            if n == n.p.r && n.getGp() != nil && n.p == n.getGp().r {
                n.p.c = true
                if n.getGp() != nil {
                    n.getGp().c = false
                }
                tree.rotateL(n.p)
            }
        }

    }
}

//test print ----------------------------------------------
var fstr = make([]string, 20, 20)

func printT(tree *node) {
    fmt.Println(tree)
    fstr = make([]string, 8, 8)
    printTree(tree, 0, "")
    for i, str := range fstr {
        fmt.Println("L", i, str)
    }
}

func printTree(tree *node, i int, n string) {
    fmt.Println("-----", tree)
    i++
    str := " "
    for n := i; n < 9; n++ {
        str += "-"
    }
    var tmp string
    if tree.c == true {
        tmp = fmt.Sprintf("%s\033[40;37m [%d] \033[0m%s%s", str, tree.v, n, str)
    } else {
        tmp = fmt.Sprintf("%s\033[41;37m [%d] \033[0m%s%s", str, tree.v, n, str)
    }

    fstr[i] += tmp
    if tree.l != nil {
        printTree(tree.l, i, "L"+fmt.Sprintf("%d", tree.v))
    }
    if tree.r != nil {
        printTree(tree.r, i, "R"+fmt.Sprintf("%d", tree.v))
    }
}

改进了下 BST insert的写法，终于轻松多了