implement rbtree delete

pkg/rbtbook/node.go

@@ -9,51 +9,6 @@ const (
 	Black = Color(true)
 )
 
-type Tree struct {
-	Root *Node
-}
-
-func (t *Tree) UpdateOrInsert(key, val fixedpoint.Value) {
-	var y *Node = nil
-	var x = t.Root
-	var node = &Node{
-		Price:  key,
-		Volume: val,
-	}
-
-	for x != nil {
-		y = x
-
-		// matched an existing node
-		if node.Price == x.Price {
-			x.Volume = val
-			return
-		} else if node.Price < x.Price {
-			x = x.Left
-		} else {
-			x = x.Right
-		}
-	}
-
-	node.Parent = y
-
-	if y == nil {
-		t.Root = node
-	} else if node.Price < y.Price {
-		y.Left = node
-	} else {
-		y.Right = node
-	}
-
-	node.Left = nil
-	node.Right = nil
-	node.Color = Red
-}
-
-func InsertFixup(tree *Tree, node *Node) {
-
-}
-
 /*
 Node
 A red node always has black children.
@@ -62,56 +17,5 @@ A black node may have red or black children
 type Node struct {
 	Left, Right, Parent *Node
 	Color               Color
-	Price               fixedpoint.Value
-	Volume              fixedpoint.Value
-}
-
-// 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 RotateLeft(tree *Tree, x *Node) {
-	var y = x.Right
-	x.Right = y.Left
-
-	if y.Left != nil {
-		y.Left.Parent = x
-	}
-
-	y.Parent = x.Parent
-
-	if x.Parent == nil {
-		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 RotateRight(tree *Tree, y *Node) {
-	x := y.Left
-	y.Left = x.Right
-
-	if x.Right != nil {
-		x.Right.Parent = y
-	}
-
-	x.Parent = y.Parent
-
-	if y.Parent == nil {
-		tree.Root = x
-	} else if y == y.Parent.Left {
-		y.Parent.Left = x
-	} else {
-		y.Parent.Right = x
-	}
-
-	x.Right = y
-	y.Parent = x
+	Key, Value          fixedpoint.Value
 }

pkg/rbtbook/tree.go

@@ -0,0 +1,294 @@
+package rbtbook
+
+import (
+	"github.com/c9s/bbgo/pkg/fixedpoint"
+)
+
+type Tree struct {
+	Root, Neel *Node
+}
+
+func NewTree() *Tree {
+	var neel = &Node{
+		Color: Black,
+	}
+	var root = neel
+	root.Parent = neel
+
+	return &Tree{
+		Root: root,
+		Neel: neel,
+	}
+}
+
+func (tree *Tree) Delete(key fixedpoint.Value) bool {
+	var del = tree.Search(key)
+	if del == nil {
+		return false
+	}
+
+	// y = the node to be deleted
+	// x (the child of the deleted node)
+	var x, y *Node
+
+	if del.Left == tree.Neel || del.Right == tree.Neel {
+		y = del
+	} else {
+		y = tree.Successor(del)
+	}
+
+	if y.Left != tree.Neel {
+		x = y.Left
+	} else {
+		x = y.Right
+	}
+
+	x.Parent = y.Parent
+
+	if y.Parent == tree.Neel {
+		tree.Root = x
+	} else if y == y.Parent.Left {
+		y.Parent.Left = x
+	} else {
+		y.Parent.Right = x
+	}
+
+	if y != del {
+		del.Key = y.Key
+	}
+
+	if y.Color == Black {
+		tree.DeleteFixup(x)
+	}
+
+	return true
+}
+
+func (tree *Tree) DeleteFixup(current *Node) {
+	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 *Tree) Insert(key, val fixedpoint.Value) {
+	var y = tree.Neel
+	var x = tree.Root
+	var node = &Node{
+		Key:   key,
+		Value: val,
+		Color: Red,
+	}
+
+	for x != tree.Neel {
+		y = x
+
+		if node.Key < x.Key {
+			x = x.Left
+		} else {
+			x = x.Right
+		}
+	}
+
+	node.Parent = y
+
+	if y == tree.Neel {
+		tree.Root = node
+	} else if node.Key < y.Key {
+		y.Left = node
+	} else {
+		y.Right = node
+	}
+
+	node.Left = tree.Neel
+	node.Right = tree.Neel
+	node.Color = Red
+
+	tree.InsertFixup(node)
+}
+
+func (tree *Tree) Search(key fixedpoint.Value) *Node {
+	var current = tree.Root
+	for current != nil && key != current.Key {
+		if key < current.Key {
+			current = current.Left
+		} else {
+			current = current.Right
+		}
+	}
+	return current
+}
+
+func (tree *Tree) InsertFixup(current *Node) {
+	// A red node can't have a red parent, we need to fix it up
+	for current.Parent.Color == Red {
+		grandParent := current.Parent.Parent
+		if current.Parent == grandParent.Left {
+			uncle := grandParent.Right
+			if uncle.Color == Red {
+				current.Parent.Color = Black
+				uncle.Color = Black
+				grandParent.Color = Red
+				current = grandParent
+			} 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
+			} 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 *Tree) RotateLeft(x *Node) {
+	var y = x.Right
+	x.Right = y.Left
+
+	if y.Left != nil {
+		y.Left.Parent = x
+	}
+
+	y.Parent = x.Parent
+
+	if x.Parent == nil {
+		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 *Tree) RotateRight(y *Node) {
+	x := y.Left
+	y.Left = x.Right
+
+	if x.Right != nil {
+		x.Right.Parent = y
+	}
+
+	x.Parent = y.Parent
+
+	if y.Parent == nil {
+		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 *Tree) LeftMost(current *Node) *Node {
+	for current.Left != nil {
+		current = current.Left
+	}
+
+	return current
+}
+
+func (tree *Tree) Successor(current *Node) *Node {
+	if current.Right != nil {
+		return tree.LeftMost(current.Right)
+	}
+
+	var newNode = current.Parent
+	for newNode != nil && current == newNode.Right {
+		current = newNode
+		newNode = newNode.Parent
+	}
+
+	return newNode
+}

pkg/rbtbook/tree_test.go

@@ -0,0 +1,37 @@
+package rbtbook
+
+import (
+	"testing"
+
+	"github.com/c9s/bbgo/pkg/fixedpoint"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestTree(t *testing.T) {
+	tree := NewTree()
+	tree.Insert(fixedpoint.NewFromFloat(3000.0), fixedpoint.NewFromFloat(10.0))
+	assert.NotNil(t, tree.Root)
+
+	tree.Insert(fixedpoint.NewFromFloat(4000.0), fixedpoint.NewFromFloat(10.0))
+	tree.Insert(fixedpoint.NewFromFloat(2000.0), fixedpoint.NewFromFloat(10.0))
+
+	// root is always black
+	assert.Equal(t, fixedpoint.NewFromFloat(3000.0), tree.Root.Key)
+	assert.Equal(t, Black, tree.Root.Color)
+
+	assert.Equal(t, fixedpoint.NewFromFloat(2000.0), tree.Root.Left.Key)
+	assert.Equal(t, Red, tree.Root.Left.Color)
+
+	assert.Equal(t, fixedpoint.NewFromFloat(4000.0), tree.Root.Right.Key)
+	assert.Equal(t, Red, tree.Root.Right.Color)
+
+	// should rotate
+	tree.Insert(fixedpoint.NewFromFloat(1500.0), fixedpoint.NewFromFloat(10.0))
+	tree.Insert(fixedpoint.NewFromFloat(1000.0), fixedpoint.NewFromFloat(10.0))
+
+	deleted := tree.Delete(fixedpoint.NewFromFloat(1000.0))
+	assert.True(t, deleted)
+
+	deleted = tree.Delete(fixedpoint.NewFromFloat(1500.0))
+	assert.True(t, deleted)
+}