红黑树插入

golang实现简单版本（没有额外的头结点存储最左最右数据）的红黑树插入

结构定义：

type RBColor bool

const (
	red   RBColor = true
	black RBColor = false
)

type RBNode struct {
	leftChild, rightChild, parent *RBNode
	value                         int
	color                         RBColor
}

type RBTree struct {
	root *RBNode // 简单起见，没有用额外的头结点
	null *RBNode // 空节点，相当于nil，每一个有数据的叶子节点的左右孩子都是它
}

实现：

func NewRBTree() *RBTree {
	return &RBTree{null: &RBNode{color: black}}
}

func (tree *RBTree) Add(newValue int) {
	if tree.root == nil {
		newNode := &RBNode{
			value:      newValue,
			color:      black,
			leftChild:  tree.null,
			rightChild: tree.null}
		tree.root = newNode
		return
	}
	node, exist := tree.find(tree.root, newValue)
	if exist {
		return
	}
	newNode := &RBNode{
		value:      newValue,
		color:      red,
		leftChild:  tree.null,
		rightChild: tree.null,
	}
	newNode.parent = node
	if newValue < node.value {
		node.leftChild = newNode
		tree.adjust(newNode)
	} else {
		node.rightChild = newNode
		tree.adjust(newNode)
	}
}

func (tree *RBTree) adjust(node *RBNode) {
	if node == tree.root { // 调整颜色向上递归到根节点了
		node.color = black
		return
	}
	parent := node.parent
	if parent.color == black { // 无需调整
		return
	}
	grandfather := parent.parent
	uncle := grandfather.leftChild
	if parent == grandfather.leftChild {
		uncle = grandfather.rightChild
	}
	if uncle.color == black { // 只旋转和调整当前三个节点的颜色，不递归向上调整
		if parent == grandfather.leftChild {
			if node == parent.leftChild { // / 形状
				tree.rightRotate(parent)
				parent.color = black
				node.color = red
				grandfather.color = red
			} else { // < 形状
				tree.leftRotate(node)
				tree.rightRotate(node)
				node.color = black
				parent.color = red
				grandfather.color = red
			}
		} else {
			if node == parent.leftChild { // > 形状
				tree.rightRotate(node)
				tree.leftRotate(node)
				node.color = black
				parent.color = red
				grandfather.color = red
			} else { // \ 形状
				tree.leftRotate(parent)
				parent.color = black
				grandfather.color = red
				node.color = red
			}
		}
	} else { // 叔叔节点是红的，只变色
		grandfather.color = red
		parent.color = black
		uncle.color = black
		tree.adjust(grandfather) // 递归向上调整
	}
}

// 只旋转，不调整颜色，让parent左旋，变成center的左子树
func (tree *RBTree) leftRotate(center *RBNode) {
	parent := center.parent
	if parent == tree.root { // 特殊情况，根节点变动
		tree.root = center
		// center一定在parent的右边
		parent.rightChild = center.leftChild
		center.leftChild = parent
		parent.parent = center
		center.parent = nil
		return
	}
	grandfather := parent.parent
	// center一定在parent的右边
	parent.rightChild = center.leftChild
	center.leftChild = parent
	parent.parent = center
	center.parent = grandfather
	if parent == grandfather.leftChild {
		grandfather.leftChild = center
	} else {
		grandfather.rightChild = center
	}
}

// 只旋转，不调整颜色，让parent右旋，变成center的右子树
func (tree *RBTree) rightRotate(center *RBNode) {
	parent := center.parent
	if parent == tree.root { // 特殊情况，根节点变动
		tree.root = center
		// center一定在parent的左边
		parent.leftChild = center.rightChild
		center.rightChild = parent
		parent.parent = center
		center.parent = nil
		return
	}
	grandfather := parent.parent
	// center一定在parent的左边
	parent.leftChild = center.rightChild
	center.rightChild = parent
	parent.parent = center
	center.parent = grandfather
	if parent == grandfather.leftChild {
		grandfather.leftChild = center
	} else {
		grandfather.rightChild = center
	}
}

// 找到返回node,true，否则返回停止的叶子节点,false
// root 一定不能为nil
func (tree *RBTree) find(root *RBNode, value int) (*RBNode, bool) {
	if root.value == value {
		return root, true
	}
	if value < root.value {
		if root.leftChild == tree.null {
			return root, false
		}
		return tree.find(root.leftChild, value)
	}
	if root.rightChild == tree.null {
		return root, false
	}
	return tree.find(root.rightChild, value)
}

// 递归遍历
func (tree *RBTree) ForEach(f func(value int, color string)) {
	if tree.root == nil {
		return
	}
	tree.innerForEach(tree.root, f)
}

func (tree *RBTree) innerForEach(root *RBNode, f func(value int, color string)) {
	if root == tree.null {
		return
	}
	tree.innerForEach(root.leftChild, f)
	color := "red"
	if root.color == black {
		color = "blk"
	}
	f(root.value, color)
	tree.innerForEach(root.rightChild, f)
}

测试：

func TestXxx(t *testing.T) {
	tree := something.NewRBTree()
	arr := []int{}
	for i := 1; i <= 20; i++ { // 生成20个数字
		arr = append(arr, i)
	}
	rand.Seed(time.Now().Unix())
	// 洗牌算法
	for i := len(arr) - 1; i >= 1; i-- {
		// 每次产生[0,i]的数字
		x := rand.Intn(i + 1)
		arr[i], arr[x] = arr[x], arr[i]
	}

	fmt.Printf("arr: %v\n", arr)
	for _, num := range arr {
		tree.Add(num)
	}
	tree.ForEach(func(i int, s string) {
		fmt.Printf("(%d, %s) ", i, s)
	})
	println()
}

结果：

用网页工具对比一下，结果是对的

总结一下插入规则：
新插入的节点都是红色。
开始调整。
如果是根节点，变黑并结束。
如果父节点是黑色，结束。
父节点是红色，两个红色相连，要调整：
1. 如果叔叔节点是红色，将父节点和叔叔节点染黑，祖父节点染红。以祖父节点继续递归向上调整。
2. 叔叔节点是黑色，以当前节点、父节点、祖父节点三者状态为准，四种形状(/ \ > <)，操作同AVL树，进行一次和两次旋转，把旋转完后的新父节点染黑，两个新子节点染红，结束。