var RED bool = true
var BLACK bool = false

type Node struct {
	Key     int     // 用于比较的键
	Value   int     // 该节点存储的值
	Left    *Node   // 左孩子
	Right   *Node   // 右孩子
	N       int     // 以当前节点为根的子树具有的节点数
	Color   bool    // 节点的颜色
}

// 返回以当前节点为根节点子树中节点总数
func (n *Node) size() int {
	if n == nil {
		return 0
	}
	return n.N
}

// 检查当前节点是否为红节点，默认空节点为黑色
func (n *Node) isRed() bool {
	if n == nil {
		return false
	}
	return n.Color == RED
}

func (n *Node) rotateLeft() *Node {
	x := n.Right
	n.Right = x.Left
	x.Left = n
	x.Color = n.Color
	n.Color = RED
	x.N = n.N // x现在的N是n之前的N
	n.N = n.Left.size() + n.Right.size() + 1
	return x
}

func (n *Node) flipColorRootRed() {
	if n == nil {
		return
	}
	n.Color = RED
	n.Left.Color = BLACK
	n.Right.Color = BLACK
}

func (n *Node) rotateRight() *Node {
	x := n.Left
	n.Left = x.Right
	x.Right = n
	x.Color = n.Color
	n.Color = RED
	x.N = n.N
	n.N = n.Left.size() + n.Right.size() + 1
	return x
}

func (n *Node) put(key, value int) *Node {
	root := n.putNode(key, value)
	root.Color = BLACK  // 保证根节点是黑色
	return root
}

func (n *Node) putNode(key, value int) *Node {
	if n == nil {  // 新插入的节点总是假定为红色
		return &Node{Key: key, Value: value, Color: RED, N: 1}
	}
	if key < n.Key {
		n.Left = n.Left.putNode(key, value)
	} else if key > n.Key {
		n.Right = n.Right.putNode(key, value)
	} else {
		n.Value = value
	}

	if n.Right.isRed() && !n.Left.isRed() {
		n = n.rotateLeft()
	}
	if n.Left.isRed() && n.Left.Left.isRed() {
		n = n.rotateRight()
	}
	if n.Left.isRed() && n.Right.isRed() {
		n.flipColorRootRed()
	}
	n.N = n.Left.size() + n.Right.size() + 1
	return n
}

// 将根节点转换成为一个临时4-节点
func (n *Node) flipColorRootBlack() {
	if n == nil {
		return
	}
	n.Color = BLACK
	n.Left.Color = RED
	n.Right.Color = RED
}

// 从右向左借
func (n *Node) moveRedLeft() *Node {
	n.flipColorRootBlack()    // 先创建一个临时4-节点（默认从父结点借）
	if n.Right.Left.isRed() { // 如果兄弟是3-节点，就从兄弟节点里面借
		n.Right = n.Right.rotateRight()
		n = n.rotateLeft()
	}
	return n
}

// 从左向右借
func (n *Node) moveRedRight() *Node {
	n.flipColorRootBlack()
	if !n.Left.Left.isRed() {  // 说明左边比右边多节点，可以借给右边
		n = n.rotateRight()
	}
	return n
}

func (n *Node) deleteMin() *Node {
	// 这是一层单独对当前根节点的包装（注意当前根节点不一定是整棵树的根节点）
	// 如果根节点的左右孩子都不是红节点（在操作过程中可能会产生临时4-节点），
	// 意味着当根节点是一个2-节点，它需要向上借节点
	if !n.Left.isRed() && !n.Right.isRed() {
		n.Color = RED
	}
	n = n.doDeleteMin()
	if n != nil {
		n.Color = BLACK
	}
	return n
}

func (n *Node) doDeleteMin() *Node {
	if n.Left == nil {
		return nil
	}
	if !n.Left.isRed() && !n.Left.Left.isRed() {
		n = n.moveRedLeft() // 如果左侧是2-节点，就需要从父结点或兄弟节点借
	}
	n.Left = n.Left.doDeleteMin()
	n = n.balance() // 在删除之后，根据红黑树的性质层层向上修正
	return n
}

// 与插入操作最后的修正方式一样
func (n *Node) balance() *Node {
	if n.Right.isRed() {
		n = n.rotateLeft()
	}
	if n.Left.isRed() && n.Left.Left.isRed() {
		n = n.rotateRight()
	}
	if n.Left.isRed() && n.Right.isRed() {
		n.flipColorRootRed()
	}
	n.N = n.Left.size() + n.Right.size() + 1
	return n
}

func (n *Node) min() *Node {
	if n.Left == nil {
		return n
	}
	return n.Left.min()
}

func (n *Node) delete(key int) *Node {
	// 依然需要对当前根节点进行包装
	if !n.Left.isRed() && !n.Right.isRed() {
		n.Color = RED
	}
	n = n.doDelete(key)
	if n != nil {
		n.Color = BLACK
	}
	return n
}

func (n *Node) doDelete(key int) *Node {
	if key < n.Key {
		if !n.Left.isRed() && !n.Left.Left.isRed() { // 左边是2-节点就向父亲或兄弟借一个
			n = n.moveRedLeft()
		}
		n.Left = n.Left.doDelete(key)
	} else {
		if n.Left.isRed() { // 这个时候要删父结点或右节点，如果当前是一个3-节点，就借它们一个
			n = n.rotateRight()
		}
		if key == n.Key && n.Right == nil {
			// 经过前面的操作，此时如果右节点为空，左节点也必为空（否则就不平衡了）
			// 当前节点就是一个红色的叶子节点，可以直接删掉
			return nil
		}
		if !n.Right.isRed() && !n.Right.Left.isRed() {
			n = n.moveRedRight() // 从左边借一个
		}
		if key == n.Key { // 用后继节点替换当前节点
			t := n.Right.min()
			n.Value = t.Value
			n.Key = t.Key
			n.Right = n.Right.deleteMin()
		} else {
			n.Right = n.Right.doDelete(key)
		}
	}
	return n.balance()
}

// 按顺序打印键
func (n *Node) print() {
	if n == nil {
		return
	}
	n.Left.print()
	fmt.Print(n.Key)
	n.Right.print()
}

func main() {
	nums := []int{1, 2, 3, 4, 5, 6, 7, 8, 9}
	root := &Node{Key: nums[0], Value: nums[0], N: 1}
	for _, n := range nums[1:] {
		root = root.put(n, n*n)
	}
	root.print()
	fmt.Println()
	root = root.delete(6)
	root.print()
}