红黑树

package main

import (
    "fmt"
    "unsafe"
)

/*
定义：
1）每个结点要么是红的，要么是黑的。
2）根结点是黑的。
3）每个叶结点（叶结点即指树尾端NIL指针或NULL结点）是黑的。
4）如果一个结点是红的，那么它的俩个儿子都是黑的。
5）对于任一结点而言，其到叶结点树尾端NIL指针的每一条路径都包含相同数目的黑结点。
*/
//红黑树(平衡二叉树）
type RBTree struct {
    index, height                 int
    isRed                         bool
    parent, leftChild, rightChild *RBTree
    ptr                           unsafe.Pointer //记录数据地址
}

func NewRBTree(index int, ptr unsafe.Pointer) *RBTree {
    return &RBTree{
        index:  index,
        isRed:  false,
        height: 1,
        ptr:    ptr,
    }
}

//判断一个节点是左孩子还是右孩子，排除根节点情况
func (root *RBTree) IsLeftChild() bool {
    if root.parent == nil {
        return false
    }
    return root.parent.leftChild == root
}

//左旋 左斜线的情况
func (h *Head) LeftRotate(root *RBTree) {
    var isParentLeft bool
    zufu := root.parent.parent
    isParentLeft = root.parent.IsLeftChild()
    root.leftChild = root.parent
    root.leftChild.parent = root
    root.leftChild.height++
    //将旋转得到的叶子节点左右孩子清空内存
    root.leftChild.leftChild = nil
    root.leftChild.rightChild = nil
    if zufu != nil {
        if isParentLeft {
            zufu.leftChild = root
        } else {
            zufu.rightChild = root
        }
        root.parent = zufu
        root.height = zufu.height + 1
    } else {
        h.ptr = root
        root.parent = nil
        root.height = 1
    }

    if root.rightChild != nil {
        root.rightChild.height--
    }
}

//右旋 右斜线的情况
func (h *Head) RightRotate(root *RBTree) {
    var isParentLeft bool
    zufu := root.parent.parent
    isParentLeft = root.parent.IsLeftChild()
    root.rightChild = root.parent
    root.rightChild.parent = root
    root.rightChild.height++
    //将旋转得到的叶子节点左右孩子清空内存
    root.rightChild.leftChild = nil
    root.rightChild.leftChild = nil
    if zufu != nil {
        if isParentLeft {
            zufu.leftChild = root
        } else {
            zufu.rightChild = root
        }
        root.parent = zufu
        root.height = zufu.height + 1
    } else {
        h.ptr = root
        root.parent = nil
        root.height = 1
    }

    if root.leftChild != nil {
        root.leftChild.height--
    }
}

func (h *Head) Insert(node *RBTree) {
    p := h.ptr
    q := p
    for p != nil {
        q = p
        if node.index == p.index {
            return
        }

        if node.index < p.index {
            p = p.leftChild
        } else {
            p = p.rightChild
        }
    }

    if q.index < node.index {
        q.rightChild = node
        node.parent = q
    }
    if q.index > node.index {
        q.leftChild = node
        node.parent = q
    }
    //插入的新节点是红色
    node.isRed = true
    node.height = q.height + 1
    //如果父节点是黑色，不需要进行调整
    if !q.isRed {
        return
    }

    //父节点是红色
    //1.叔叔节点为红色 ==>父节点和叔叔节点改黑色
    if q.parent.leftChild != nil && q.parent.rightChild != nil {
        q.parent.leftChild.isRed = false
        q.parent.rightChild.isRed = false
        return
    }

    //2.叔叔节点为空，新节点和父节点，祖父节点一条线上
    //a.父节点和新节点都是左孩子
    if q.IsLeftChild() && node.IsLeftChild() {
        h.RightRotate(q)
        q.isRed = false
        q.leftChild.isRed = true
        q.rightChild.isRed = true
        return
    }
    //b.父节点和新节点都是右孩子
    if !q.IsLeftChild() && !node.IsLeftChild() {
        h.LeftRotate(q)
        q.isRed = false
        q.leftChild.isRed = true
        q.rightChild.isRed = true
        return
    }

    //3.叔叔节点为空，新节点和父节点，祖父节点不在一条线上
    //a.父节点为右孩子，新节点为左孩子
    if !q.IsLeftChild() && node.IsLeftChild() {
        //对新节点右旋,变成一条线上
        h.RightRotate(node)
        //再对新节点左旋
        h.LeftRotate(node)
        node.isRed = false
        node.leftChild.isRed = true
        node.rightChild.isRed = true
        return
    }
    //b.父节点为左孩子，新节点为右孩子
    if q.IsLeftChild() && !node.IsLeftChild() {
        //对新节点左旋，变成一条线上
        h.LeftRotate(node)
        //再对新节点右旋
        h.RightRotate(node)
        node.isRed = false
        node.leftChild.isRed = true
        node.rightChild.isRed = true
        return
    }
}

//层次遍历
func Print(root *RBTree) {
    queq := make([]*RBTree, 0)
    if root == nil {
        return
    }
    i := 0
    queq = append(queq, root)
    for i < len(queq) {
        fmt.Println(queq[i].index, queq[i].height, queq[i].isRed, queq[i].parent)
        if queq[i].leftChild != nil {
            queq = append(queq, queq[i].leftChild)
        }
        if queq[i].rightChild != nil {
            queq = append(queq, queq[i].rightChild)
        }
        i++
    }
}

type Data struct {
    id   int
    name string
}

type Head struct {
    ptr *RBTree
}

func main() {
    //用链表头指向树的root地址
    test := []int{8, 7, 4, 5, 3, 2, 1}
    root := NewRBTree(6, unsafe.Pointer(&Data{10, "xxx"}))
    head := Head{root}
    fmt.Println(&root)
    fmt.Println(unsafe.Pointer(&root))
    for _, v := range test {
        head.Insert(&RBTree{index: v})
    }
    fmt.Println(&root)
    Print(head.ptr)
}