redblack/redblack.go

package redblack

type Comparable interface {
	Compare(c Comparable) int
}

type Tree struct {
	root *Node
	num  int
}

type Node struct {
	key    Comparable
	color  int
	parent *Node
	left   *Node
	right  *Node
}

func (t *Tree) rightRotate(node *Node) {
	var (
		left = node.left
	)

	node.left = left.right
	if node.left != nil {
		node.left.parent = node
	}
	left.parent = node.parent

	if node.parent == nil {
		t.root = left
	} else if node == node.parent.left {
		node.parent.left = left
	} else {
		node.parent.right = left
	}

	left.right = node
	node.parent = left
}

func (t *Tree) leftRotate(node *Node) {
	var (
		right = node.right
	)

	node.right = right.left
	if node.right != nil {
		node.right.parent = node
	}
	right.parent = node.parent

	if node.parent == nil {
		t.root = right
	} else if node == node.parent.left {
		node.parent.left = right
	} else {
		node.parent.right = right
	}

	right.left = node
	node.parent = right
}

func (t *Tree) bst(trav, temp *Node) *Node {
	if trav == nil {
		return temp
	}

	if temp.key.Compare(trav.key) < 0 {
		trav.left = t.bst(trav.left, temp)
		trav.left.parent = trav
	} else if temp.key.Compare(trav.key) > 0 {
		trav.right = t.bst(trav.right, temp)
		trav.right.parent = trav
	} else {
		panic("duplicate key") // duplicate key not allowed
	}

	return trav
}

func (t *Tree) fixup(root, pt *Node) {
	var (
		parentPt      *Node
		grandParentPt *Node
		unclePt       *Node
	)

	for pt != root && pt.color != 0 && pt.parent.color == 1 {
		parentPt = pt.parent
		grandParentPt = pt.parent.parent

		//case A: Parent of pt is left child of grand-parent of pt
		if parentPt == grandParentPt.left {
			unclePt = grandParentPt.right

			//case 1: the uncle of pt is also red; only recoloring required
			if unclePt != nil && unclePt.color == 1 {
				grandParentPt.color = 1
				parentPt.color = 0
				unclePt.color = 0
				pt = grandParentPt
			} else {
				//case 2: pt is right child of its parent; left-rotation required
				if pt == parentPt.right {
					t.leftRotate(parentPt)
					pt = parentPt
					parentPt = pt.parent
				}

				//case 3: pt is left child of its parent; right-rotation required
				t.rightRotate(grandParentPt)
				t := parentPt.color
				parentPt.color = grandParentPt.color
				grandParentPt.color = t
				pt = parentPt
			}
		} else {
			//case B: parent of pt is right child of grand parent of pt
			unclePt = grandParentPt.left

			//case 1: the unclePt is also red; only recoloring required
			if unclePt != nil && unclePt.color == 1 {
				grandParentPt.color = 1
				parentPt.color = 0
				unclePt.color = 0
				pt = grandParentPt
			} else {

				//case 2: pt is left child of its parent; right-rotate required
				if pt == parentPt.left {
					t.rightRotate(parentPt)
					pt = parentPt
					parentPt = pt.parent
				}

				//case 3: pt is right child of its parent; left-rotate required
				t.leftRotate(grandParentPt)
				t := parentPt.color
				parentPt.color = grandParentPt.color
				grandParentPt.color = t
				pt = parentPt
			}
		}
	}
}

func (t *Tree) Insert(key Comparable) {
	temp := &Node{
		key:   key,
		color: 1,
	}

	t.root = t.bst(t.root, temp)
	t.fixup(t.root, temp)
	t.root.color = 0
	t.num++
}

func (t *Tree) Slice() []Comparable {
	returnValue := make([]Comparable, t.num)
	t.slice(returnValue, 0, t.root)
	return returnValue
}

func (t *Tree) slice(s []Comparable, index int, node *Node) int {
	if node.left != nil {
		index = t.slice(s, index, node.left)
	}

	s[index] = node.key
	index++

	if node.right != nil {
		index = t.slice(s, index, node.right)
	}

	return index
}