bbgo/pkg/types/rbtree.go

package types

import (
	"fmt"

	"git.qtrade.icu/lychiyu/bbgo/pkg/fixedpoint"
)

type RBTree struct {
	Root *RBNode
	size int
}

func NewRBTree() *RBTree {
	var root = NewNil()
	root.parent = NewNil()
	return &RBTree{
		Root: root,
	}
}

func (tree *RBTree) Delete(key fixedpoint.Value) bool {
	var deleting = tree.Search(key)
	if deleting == nil {
		return false
	}

	// y = the node to be deleted
	// x (the child of the deleted node)
	var x, y *RBNode
	// fmt.Printf("neel = %p %+v\n", neel, neel)
	// fmt.Printf("deleting = %+v\n", deleting)

	// the deleting node has only one child, it's easy,
	// we just connect the child the parent of the deleting node
	if deleting.left.isNil() || deleting.right.isNil() {
		y = deleting
		// fmt.Printf("y = deleting = %+v\n", y)
	} else {
		// if both children are not NIL (neel), we need to find the successor
		// and copy the successor to the memory location of the deleting node.
		// since it's successor, it always has no child connecting to it.
		y = tree.Successor(deleting)
		// fmt.Printf("y = successor = %+v\n", y)
	}

	// y.left or y.right could be neel
	if y.left.isNil() {
		x = y.right
	} else {
		x = y.left
	}

	// fmt.Printf("x = %+v\n", y)
	x.parent = y.parent

	if y.parent.isNil() {
		tree.Root = x
	} else if y == y.parent.left {
		y.parent.left = x
	} else {
		y.parent.right = x
	}

	// copy the data from the successor to the memory location of the deleting node
	if y != deleting {
		deleting.key = y.key
		deleting.value = y.value
	}

	if y.color == Black {
		tree.DeleteFixup(x)
	}

	tree.size--

	return true
}

func (tree *RBTree) DeleteFixup(current *RBNode) {
	for current != tree.Root && current.color == Black {
		if current == current.parent.left {
			sibling := current.parent.right
			if sibling.color == Red {
				sibling.color = Black
				current.parent.color = Red
				tree.RotateLeft(current.parent)
				sibling = current.parent.right
			}

			// if both are black nodes
			if sibling.left.color == Black && sibling.right.color == Black {
				sibling.color = Red
				current = current.parent
			} else {
				// only one of the child is black
				if sibling.right.color == Black {
					sibling.left.color = Black
					sibling.color = Red
					tree.RotateRight(sibling)
					sibling = current.parent.right
				}

				sibling.color = current.parent.color
				current.parent.color = Black
				sibling.right.color = Black
				tree.RotateLeft(current.parent)
				current = tree.Root
			}
		} else { // if current is right child
			sibling := current.parent.left
			if sibling.color == Red {
				sibling.color = Black
				current.parent.color = Red
				tree.RotateRight(current.parent)
				sibling = current.parent.left
			}

			if sibling.left.color == Black && sibling.right.color == Black {
				sibling.color = Red
				current = current.parent
			} else { // if only one of child is Black

				// the left child of sibling is black, and right child is red
				if sibling.left.color == Black {
					sibling.right.color = Black
					sibling.color = Red
					tree.RotateLeft(sibling)
					sibling = current.parent.left
				}

				sibling.color = current.parent.color
				current.parent.color = Black
				sibling.left.color = Black
				tree.RotateRight(current.parent)
				current = tree.Root
			}
		}
	}

	current.color = Black
}

func (tree *RBTree) Upsert(key, val fixedpoint.Value) {
	var y = NewNil()
	var x = tree.Root
	var node = &RBNode{
		key:    key,
		value:  val,
		color:  Red,
		left:   NewNil(),
		right:  NewNil(),
		parent: NewNil(),
	}

	for !x.isNil() {
		y = x

		if node.key == x.key {
			// found node, skip insert and fix
			x.value = val
			return
		} else if node.key.Compare(x.key) < 0 {
			x = x.left
		} else {
			x = x.right
		}
	}

	node.parent = y

	if y.isNil() {
		tree.Root = node
	} else if node.key.Compare(y.key) < 0 {
		y.left = node
	} else {
		y.right = node
	}

	tree.InsertFixup(node)
}

func (tree *RBTree) Insert(key, val fixedpoint.Value) {
	var y = NewNil()
	var x = tree.Root
	var node = &RBNode{
		key:    key,
		value:  val,
		color:  Red,
		left:   NewNil(),
		right:  NewNil(),
		parent: NewNil(),
	}

	for !x.isNil() {
		y = x

		if node.key.Compare(x.key) < 0 {
			x = x.left
		} else {
			x = x.right
		}
	}

	node.parent = y

	if y.isNil() {
		tree.Root = node
	} else if node.key.Compare(y.key) < 0 {
		y.left = node
	} else {
		y.right = node
	}

	tree.size++
	tree.InsertFixup(node)
}

func (tree *RBTree) Search(key fixedpoint.Value) *RBNode {
	var current = tree.Root
	for !current.isNil() && key != current.key {
		if key.Compare(current.key) < 0 {
			current = current.left
		} else {
			current = current.right
		}
	}

	if current.isNil() {
		return nil
	}

	return current
}

func (tree *RBTree) Size() int {
	return tree.size
}

func (tree *RBTree) InsertFixup(current *RBNode) {
	// A red node can't have a red parent, we need to fix it up
	for current.parent.color == Red {
		if current.parent == current.parent.parent.left {
			uncle := current.parent.parent.right
			if uncle.color == Red {
				current.parent.color = Black
				uncle.color = Black
				current.parent.parent.color = Red
				current = current.parent.parent
			} else { // if uncle is black
				if current == current.parent.right {
					current = current.parent
					tree.RotateLeft(current)
				}

				current.parent.color = Black
				current.parent.parent.color = Red
				tree.RotateRight(current.parent.parent)
			}
		} else {
			uncle := current.parent.parent.left
			if uncle.color == Red {
				current.parent.color = Black
				uncle.color = Black
				current.parent.parent.color = Red
				current = current.parent.parent
			} else {
				if current == current.parent.left {
					current = current.parent
					tree.RotateRight(current)
				}

				current.parent.color = Black
				current.parent.parent.color = Red
				tree.RotateLeft(current.parent.parent)
			}
		}
	}

	// ensure that root is black
	tree.Root.color = Black
}

// RotateLeft
// x is the axes of rotation, y is the node that will be replace x's position.
// we need to:
// 1. move y's left child to the x's right child
// 2. change y's parent to x's parent
// 3. change x's parent to y
func (tree *RBTree) RotateLeft(x *RBNode) {
	var y = x.right
	x.right = y.left

	if !y.left.isNil() {
		y.left.parent = x
	}

	y.parent = x.parent

	if x.parent.isNil() {
		tree.Root = y
	} else if x == x.parent.left {
		x.parent.left = y
	} else {
		x.parent.right = y
	}

	y.left = x
	x.parent = y
}

func (tree *RBTree) RotateRight(y *RBNode) {
	x := y.left
	y.left = x.right

	if !x.right.isNil() {
		if x.right == nil {
			panic(fmt.Errorf("x.right is nil: node = %+v, left = %+v, right = %+v, parent = %+v", x, x.left, x.right, x.parent))
		}
		x.right.parent = y
	}

	x.parent = y.parent

	if y.parent.isNil() {
		tree.Root = x
	} else if y == y.parent.left {
		y.parent.left = x
	} else {
		y.parent.right = x
	}

	x.right = y
	y.parent = x
}

func (tree *RBTree) Rightmost() *RBNode {
	return tree.RightmostOf(tree.Root)
}

func (tree *RBTree) RightmostOf(current *RBNode) *RBNode {
	if current.isNil() || current == nil {
		return nil
	}

	for !current.right.isNil() {
		current = current.right
	}

	return current
}

func (tree *RBTree) Leftmost() *RBNode {
	return tree.LeftmostOf(tree.Root)
}

func (tree *RBTree) LeftmostOf(current *RBNode) *RBNode {
	if current.isNil() || current == nil {
		return nil
	}

	for !current.left.isNil() {
		current = current.left
	}

	return current
}

func (tree *RBTree) Successor(current *RBNode) *RBNode {
	if !current.right.isNil() {
		return tree.LeftmostOf(current.right)
	}

	var newNode = current.parent
	for !newNode.isNil() && current == newNode.right {
		current = newNode
		newNode = newNode.parent
	}

	return newNode
}

func (tree *RBTree) Preorder(cb func(n *RBNode)) {
	tree.PreorderOf(tree.Root, cb)
}

func (tree *RBTree) PreorderOf(current *RBNode, cb func(n *RBNode)) {
	if !current.isNil() && current != nil {
		cb(current)
		tree.PreorderOf(current.left, cb)
		tree.PreorderOf(current.right, cb)
	}
}

// Inorder traverses the tree in ascending order
func (tree *RBTree) Inorder(cb func(n *RBNode) bool) {
	tree.InorderOf(tree.Root, cb)
}

func (tree *RBTree) InorderOf(current *RBNode, cb func(n *RBNode) bool) {
	if !current.isNil() && current != nil {
		tree.InorderOf(current.left, cb)
		if !cb(current) {
			return
		}
		tree.InorderOf(current.right, cb)
	}
}

// InorderReverse traverses the tree in descending order
func (tree *RBTree) InorderReverse(cb func(n *RBNode) bool) {
	tree.InorderReverseOf(tree.Root, cb)
}

func (tree *RBTree) InorderReverseOf(current *RBNode, cb func(n *RBNode) bool) {
	if !current.isNil() && current != nil {
		tree.InorderReverseOf(current.right, cb)
		if !cb(current) {
			return
		}
		tree.InorderReverseOf(current.left, cb)
	}
}

func (tree *RBTree) Postorder(cb func(n *RBNode) bool) {
	tree.PostorderOf(tree.Root, cb)
}

func (tree *RBTree) PostorderOf(current *RBNode, cb func(n *RBNode) bool) {
	if !current.isNil() && current != nil {
		tree.PostorderOf(current.left, cb)
		tree.PostorderOf(current.right, cb)
		if !cb(current) {
			return
		}
	}
}

func (tree *RBTree) CopyInorderReverse(limit int) *RBTree {
	newTree := NewRBTree()
	if limit == 0 {
		tree.InorderReverse(copyNodeFast(newTree))
		return newTree
	}

	tree.InorderReverse(copyNodeLimit(newTree, limit))
	return newTree
}

func (tree *RBTree) CopyInorder(limit int) *RBTree {
	newTree := NewRBTree()
	if limit == 0 {
		tree.Inorder(copyNodeFast(newTree))
		return newTree
	}

	tree.Inorder(copyNodeLimit(newTree, limit))
	return newTree
}

func (tree *RBTree) Print() {
	tree.Inorder(func(n *RBNode) bool {
		fmt.Printf("%v -> %v\n", n.key, n.value)
		return true
	})
}

func copyNodeFast(newTree *RBTree) func(n *RBNode) bool {
	return func(n *RBNode) bool {
		newTree.Insert(n.key, n.value)
		return true
	}
}

func copyNodeLimit(newTree *RBTree, limit int) func(n *RBNode) bool {
	cnt := 0
	return func(n *RBNode) bool {
		if limit > 0 && cnt >= limit {
			return false
		}

		newTree.Insert(n.key, n.value)
		cnt++
		return true
	}
}