Deletion from a Red-Black Tree

package main

import "fmt"

type Color bool

const (
	RED   Color = true
	BLACK Color = false
)

type Node struct {
	data   int
	color  Color
	left   *Node
	right  *Node
	parent *Node
}

type RedBlackTree struct {
	TNULL *Node
	root  *Node
}

func NewRedBlackTree() *RedBlackTree {
	tnull := &Node{color: BLACK}
	return &RedBlackTree{
		TNULL: tnull,
		root:  tnull,
	}
}

func (tree *RedBlackTree) leftRotate(x *Node) {
	y := x.right
	x.right = y.left
	if y.left != tree.TNULL {
		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 *RedBlackTree) rightRotate(x *Node) {
	y := x.left
	x.left = y.right
	if y.right != tree.TNULL {
		y.right.parent = x
	}
	y.parent = x.parent
	if x.parent == nil {
		tree.root = y
	} else if x == x.parent.right {
		x.parent.right = y
	} else {
		x.parent.left = y
	}
	y.right = x
	x.parent = y
}

func (tree *RedBlackTree) insertFix(k *Node) {
	for k.parent != nil && k.parent.color == RED {
		if k.parent == k.parent.parent.right {
			u := k.parent.parent.left
			if u.color == RED {
				u.color = BLACK
				k.parent.color = BLACK
				k.parent.parent.color = RED
				k = k.parent.parent
			} else {
				if k == k.parent.left {
					k = k.parent
					tree.rightRotate(k)
				}
				k.parent.color = BLACK
				k.parent.parent.color = RED
				tree.leftRotate(k.parent.parent)
			}
		} else {
			u := k.parent.parent.right
			if u.color == RED {
				u.color = BLACK
				k.parent.color = BLACK
				k.parent.parent.color = RED
				k = k.parent.parent
			} else {
				if k == k.parent.right {
					k = k.parent
					tree.leftRotate(k)
				}
				k.parent.color = BLACK
				k.parent.parent.color = RED
				tree.rightRotate(k.parent.parent)
			}
		}
		if k == tree.root {
			break
		}
	}
	tree.root.color = BLACK
}

func (tree *RedBlackTree) Insert(key int) {
	node := &Node{
		data:   key,
		color:  RED,
		left:   tree.TNULL,
		right:  tree.TNULL,
		parent: nil,
	}

	var y *Node = nil
	x := tree.root

	for x != tree.TNULL {
		y = x
		if node.data < x.data {
			x = x.left
		} else {
			x = x.right
		}
	}

	node.parent = y
	if y == nil {
		tree.root = node
	} else if node.data < y.data {
		y.left = node
	} else {
		y.right = node
	}

	if node.parent == nil {
		node.color = BLACK
		return
	}

	if node.parent.parent == nil {
		return
	}

	tree.insertFix(node)
}

func (tree *RedBlackTree) rbTransplant(u, v *Node) {
	if u.parent == nil {
		tree.root = v
	} else if u == u.parent.left {
		u.parent.left = v
	} else {
		u.parent.right = v
	}
	v.parent = u.parent
}

func (tree *RedBlackTree) minimum(node *Node) *Node {
	for node.left != tree.TNULL {
		node = node.left
	}
	return node
}

func (tree *RedBlackTree) fixDelete(x *Node) {
	for x != tree.root && x.color == BLACK {
		if x == x.parent.left {
			s := x.parent.right
			if s.color == RED {
				s.color = BLACK
				x.parent.color = RED
				tree.leftRotate(x.parent)
				s = x.parent.right
			}

			if s.left.color == BLACK && s.right.color == BLACK {
				s.color = RED
				x = x.parent
			} else {
				if s.right.color == BLACK {
					s.left.color = BLACK
					s.color = RED
					tree.rightRotate(s)
					s = x.parent.right
				}

				s.color = x.parent.color
				x.parent.color = BLACK
				s.right.color = BLACK
				tree.leftRotate(x.parent)
				x = tree.root
			}
		} else {
			s := x.parent.left
			if s.color == RED {
				s.color = BLACK
				x.parent.color = RED
				tree.rightRotate(x.parent)
				s = x.parent.left
			}

			if s.right.color == BLACK && s.left.color == BLACK {
				s.color = RED
				x = x.parent
			} else {
				if s.left.color == BLACK {
					s.right.color = BLACK
					s.color = RED
					tree.leftRotate(s)
					s = x.parent.left
				}

				s.color = x.parent.color
				x.parent.color = BLACK
				s.left.color = BLACK
				tree.rightRotate(x.parent)
				x = tree.root
			}
		}
	}
	x.color = BLACK
}

func (tree *RedBlackTree) Delete(data int) {
	tree.deleteNodeHelper(tree.root, data)
}

func (tree *RedBlackTree) deleteNodeHelper(node *Node, key int) {
	var z *Node = tree.TNULL
	var x, y *Node

	for node != tree.TNULL {
		if node.data == key {
			z = node
		}
		if node.data <= key {
			node = node.right
		} else {
			node = node.left
		}
	}

	if z == tree.TNULL {
		fmt.Printf("Kalit %d daraxtda topilmadi\n", key)
		return
	}

	y = z
	yOriginalColor := y.color
	if z.left == tree.TNULL {
		x = z.right
		tree.rbTransplant(z, z.right)
	} else if z.right == tree.TNULL {
		x = z.left
		tree.rbTransplant(z, z.left)
	} else {
		y = tree.minimum(z.right)
		yOriginalColor = y.color
		x = y.right
		if y.parent == z {
			x.parent = y
		} else {
			tree.rbTransplant(y, y.right)
			y.right = z.right
			y.right.parent = y
		}

		tree.rbTransplant(z, y)
		y.left = z.left
		y.left.parent = y
		y.color = z.color
	}
	if yOriginalColor == BLACK {
		tree.fixDelete(x)
	}
}

func (tree *RedBlackTree) printHelper(root *Node, indent string, last bool) {
	if root != tree.TNULL {
		fmt.Print(indent)
		if last {
			fmt.Print("R----")
			indent += "   "
		} else {
			fmt.Print("L----")
			indent += "|  "
		}
		sColor := "RED"
		if root.color == BLACK {
			sColor = "BLACK"
		}
		fmt.Printf("%d(%s)\n", root.data, sColor)
		tree.printHelper(root.left, indent, false)
		tree.printHelper(root.right, indent, true)
	}
}

func (tree *RedBlackTree) PrintTree() {
	tree.printHelper(tree.root, "", true)
}

func main() {
	tree := NewRedBlackTree()
	elements := []int{55, 40, 65, 60, 75, 57}
	for _, e := range elements {
		tree.Insert(e)
	}

	fmt.Println("Boshlang'ich daraxt:")
	tree.PrintTree()

	tree.Delete(40)
	fmt.Println("\n40 o'chirilgandan so'ng:")
	tree.PrintTree()
}