red_black_tree.go

package main
import(
	"fmt"
	"strconv"
)

type Node struct{
  value int
  right *Node
  left *Node
  parent *Node
  red int
}

func(self *Node) String() string{
	description := "Value: " + strconv.Itoa(self.value) + "\n"
	description += "Color: " + self.color() + "\n"
	description += self.right_child_description() + "\n"
	description += self.left_child_description() + "\n"
	description += self.parent_description() + "\n"
	return description
}

func(self *Node) color() string{
	if isRed(self){
		return "red"
	}
	return "black"
}

func(self *Node) right_child_description() string{
	if(self.right == nil){
		return "Right child: NIL"
	}
	return "Right child: "+ strconv.Itoa(self.right.value)
}

func(self *Node) left_child_description() string{
	if(self.left == nil){
		return "Left child: NIL"
	}
	return "Left child: "+ strconv.Itoa(self.left.value)
}

func(self *Node) parent_description() string{
	if(self.parent == nil){
		return "I'm the root"
	}
	return "Parent: "+ strconv.Itoa(self.parent.value)
}

func(self *Node) height() int{
	if self.parent == nil{
		return 1
	}else{
		return self.parent.height() + 1
	}
}


func (node *Node) uncle() *Node{
	if node.grandParent() != nil{
		if node.grandParent().right == node.parent{
			return node.grandParent().left
		}else{
			return node.grandParent().right
		}
	}
	return nil
}

func (node *Node) grandParent() *Node{
	if node.parent != nil{
		return node.parent.parent;
	}else{
		return nil
	}
}

func isRed(node *Node) bool{
	if node != nil && node.red == 1{
		return true
	}
	return false
}


type Tree struct{
	root *Node
}

func (self *Tree) add(value int) {
	if self.root == nil{
		self.root = &Node{value: value, red: 0}
	}else{
		self.recursive_add(value, self.root)
	}
}

func (self *Tree) recursive_add(value int, node *Node){
	if value < node.value{
		if node.left == nil{
			node.left = &Node{value: value, parent: node, red: 1}
			self.insert_case_one(node.left)
		}else{
			self.recursive_add(value, node.left)
		}
	}else{
		if node.right == nil{
			node.right = &Node{value: value, parent: node, red: 1}
			self.insert_case_one(node.right)
		}else{
			self.recursive_add(value,node.right)
		}
	}
}

func (self *Tree) insert_case_one(node *Node){
	if node.parent == nil{
		node.red = 0
	}else{
		self.insert_case_two(node)
	}
}

func(self *Tree) insert_case_two(node *Node){
	if isRed(node.parent){
		self.insert_case_three(node)
	}
}

func (self *Tree) insert_case_three(node *Node) {
	
	uncle := node.uncle()
	parent := node.parent

	if uncle != nil && ( isRed(uncle) && isRed(parent) ){
		uncle.red = 0
		parent.red = 0
		g := node.grandParent()
		g.red = 1
		self.insert_case_one(g)
	}else{
		self.insert_case_four(node)
	}
}

func (self *Tree) insert_case_four(node *Node){
	
	g := node.grandParent()

	if(node == node.parent.right && g.left == node.parent){
		self.rotate_left(node.parent)
		node = node.left
	}else if node == node.parent.left && g.right == node.parent {
		self.rotate_right(node.parent)
		node = node.right
	}
	self.insert_case_five(node)
}

func (self *Tree) insert_case_five(node *Node){
	
	g := node.grandParent()
	node.parent.red = 0
	g.red = 1
	if node == node.parent.left{
		self.rotate_right(g)			
	}else{
		self.rotate_left(g)		
	}
}

func (self *Tree) rotate_left(node *Node){
	y := node.right 
	node.right = y.left
	if y.left != nil{
		y.left.parent = node
	}
	y.parent = node.parent
	if node.parent == nil{
		self.root = y
	}else if node == node.parent.left{
		node.parent.left = y
	}else{
		node.parent.right = y
	}
	y.left = node
	node.parent = y

}

func (self *Tree) rotate_right(node *Node){
	y := node.left 
	node.left = y.right
	if y.right != nil{
		y.right.parent = node
	}
	y.parent = node.parent
	if node.parent == nil{
		self.root = y
	}else if node == node.parent.right{
		node.parent.right = y
	}else{
		node.parent.left = y
	}
	y.right = node
	node.parent = y
}

func(self *Tree) find(value int) *Node{
	if self.root != nil{
		return self.recursive_find(value, self.root)
	}else{
		return nil
	}
}

func(self *Tree) recursive_find(value int, node *Node) *Node{
	if node.value == value {
		return node
	}else if value < node.value && node.left != nil {
		return self.recursive_find(value,node.left)
	}else if node.right != nil{
		return self.recursive_find(value,node.right)
	}else{
		return nil
	}
}


func main() {
	/* Test */
	tree := Tree{}
	tree.add(11)
	tree.add(2)
	tree.add(14)
	tree.add(1)
	tree.add(7)
	tree.add(5)
	tree.add(8)
	tree.add(15)
	tree.add(4)
	tree.add(20)
	tree.add(9)
	tree.add(10)
	fmt.Println(tree.find(2))
}