红黑树之Golang代码简释

写在前面：
一，本文侧重诠释对算法的思考记录过程，忽略其他诸如代码简洁、字符编码等细节问题。
二，本文结合 红黑树插入过程图示 这篇一起看，有助于理解。

package main

import "fmt"

type Node struct {
	Data   int
	Left   *Node
	Right  *Node
	Color  int
	Parent *Node
}

//红黑树（Red-Black Tree）是每个节点都带有颜色属性的二叉排序（查找）树，具备以下特性：
//1,节点是红色或黑色
//2,根节点是黑色的
//3,每个叶子节点都是黑色的空节点（NIL），也就是说，叶子节点不存储数据
//4,任何相邻的节点都不能同时为红色，也就是说，红色节点是被黑色节点隔开的
//5,每个节点，从该节点到达其可达叶子节点的所有路径，都包含相同数目的黑色节点
type RbTree struct {
	Tree *Node
}

const (
	RED   = 1
	BLACK = 0
)

//中序遍历
func (t *RbTree) MidOrderTraverse(tree *Node) {
	if tree == nil {
		return
	}

	if tree.Left != nil {
		t.MidOrderTraverse(tree.Left)
	}

	fmt.Println(tree)

	if tree.Right != nil {
		t.MidOrderTraverse(tree.Right)
	}
}

//插入节点
func (t *RbTree) Insert(data int) {
	//如果是空树，就插入到根节点
	if t.Tree == nil {
		t.Tree = &Node{Data: data}
		return
	}

	tree := t.Tree
	for tree != nil {
		if data < tree.Data {
			if tree.Left == nil {
				tree.Left = &Node{Data: data, Color: RED, Parent: tree}
				t.ChangeColor(tree.Left)
				return
			}
			tree = tree.Left
		} else if data > tree.Data {
			if tree.Right == nil {
				tree.Right = &Node{Data: data, Color: RED, Parent: tree}
				t.ChangeColor(tree.Right)
				return
			}
			tree = tree.Right
		}
	}

}

func (t *RbTree) ChangeColor(tree *Node) {
	flag := false //被关注节点和其父节点是否在同一侧
	//被关注节点的父节点不是根节点
	if tree.Parent.Parent != nil {
		isLeft := t.IsLeft(tree)
		parentLeft := t.IsLeft(tree.Parent)
		uncle := t.GetUncle(tree)
		if isLeft && parentLeft { //被关注节点和其父节点同在左侧
			flag = true
		} else if !isLeft && !parentLeft { //被关注节点和其父节点同在右侧
			flag = true
		}
		//case1：被关注节点的父节点和叔叔/伯伯节点都是红色
		if uncle != nil && tree.Parent.Color == RED && uncle.Color == RED {
			tree.Parent.Color = BLACK
			uncle.Color = BLACK
			if tree.Parent.Parent.Parent != nil { //祖父不是根节点
				tree.Parent.Parent.Color = RED
				tree = tree.Parent.Parent //关注节点变成祖父节点
				if tree.Parent.Parent != nil {
					isLeft = t.IsLeft(tree)
					parentLeft = t.IsLeft(tree.Parent)
					uncle = t.GetUncle(tree)
					if isLeft && parentLeft {
						flag = true
					} else if !isLeft && !parentLeft {
						flag = true
					}
				}
			}
		}

		//case2：关注节点和其父节点都是红色，但不在同一侧（同为左侧或同为右侧），且叔叔/伯伯节点是黑色
		if tree.Parent.Parent != nil && tree.Parent.Color == RED && (uncle == nil || uncle.Color == BLACK) && !flag {
			if parentLeft { //父节点在左侧
				tree = t.LeftRotate(tree.Parent).Left
				isLeft = true
				flag = true
			} else {
				tree = t.RightRotate(tree.Parent).Right
				isLeft = false
				flag = true
			}
		}

		//case3：关注节点和其父节点都是红色，且在同一侧（同为左侧或同为右侧），且叔叔/伯伯节点是黑色
		if tree.Parent.Parent != nil && tree.Parent.Color == RED && (uncle == nil || uncle.Color == BLACK) && flag {
			tree.Parent.Color = BLACK
			tree.Parent.Parent.Color = RED
			if parentLeft { //父节点在左侧
				tree = t.RightRotate(tree.Parent.Parent)
			} else {
				tree = t.LeftRotate(tree.Parent.Parent)
			}
		} else if tree.Parent.Parent != nil && uncle != nil && tree.Parent.Color == RED && uncle.Color == RED { //case1
			t.ChangeColor(tree)
		}

	}
}

func (t *RbTree) IsLeft(tree *Node) bool {
	isLeft := false
	if tree.Parent != nil && tree.Parent.Left == tree {
		isLeft = true
	}
	return isLeft
}

func (t *RbTree) GetUncle(tree *Node) *Node {
	var uncle *Node
	if tree.Parent.Parent != nil {
		uncle = tree.Parent.Parent.Left
		if tree.Parent.Parent.Left == tree.Parent {
			uncle = tree.Parent.Parent.Right
		}
	}

	return uncle
}

//右旋
func (t *RbTree) RightRotate(tree *Node) *Node {
	subTree := tree.Left
	isLeft := false
	if tree.Parent != nil {
		subTree.Parent = tree.Parent //更新新子树的父节点
		if tree.Parent.Left == tree {
			isLeft = true
		}
	} else {
		subTree.Parent = nil
	}
	tree.Left = subTree.Right //原来左节点的右子树挂到老的根节点的左子树
	if subTree.Right != nil {
		subTree.Right.Parent = tree
	}
	tree.Parent = subTree //原来的左节点变成老的根节点的父节点
	subTree.Right = tree  //原来的根节点变成原来左节点的右子树
	tree = subTree
	if tree.Parent == nil { //旋转的是整棵树的根节点
		t.Tree = tree
	} else {
		if isLeft { //更新老的子树根节点父节点指针指向新的根节点
			tree.Parent.Left = tree
		} else {
			tree.Parent.Right = tree
		}
	}
	return tree
}

//左旋
func (t *RbTree) LeftRotate(tree *Node) *Node {
	subTree := tree.Right
	isLeft := false
	if tree.Parent != nil {
		subTree.Parent = tree.Parent
		if tree.Parent.Left == tree {
			isLeft = true
		}
	} else {
		subTree.Parent = nil
	}
	tree.Right = subTree.Left
	if subTree.Left != nil {
		subTree.Left.Parent = tree
	}
	tree.Parent = subTree
	subTree.Left = tree
	tree = subTree
	if tree.Parent == nil {
		t.Tree = tree
	} else {
		if isLeft {
			tree.Parent.Left = tree
		} else {
			tree.Parent.Right = tree
		}
	}
	return tree
}

//测试数据
func main() {
	treeObj := &RbTree{}
	treeObj.Insert(99)
	// treeObj.Insert(78)
	// treeObj.Insert(120)
	// treeObj.Insert(66)
	// treeObj.Insert(83)
	// treeObj.Insert(57)
	// treeObj.Insert(52)
	// treeObj.Insert(61)
	// treeObj.Insert(64)
	// treeObj.Insert(59)
	// treeObj.Insert(60)
	// treeObj.Insert(70)
	// treeObj.Insert(75)
	// treeObj.Insert(72)
	// treeObj.Insert(74)

	treeObj.MidOrderTraverse(treeObj.Tree)
}