GO - AVL TREES CODE

JAVA PYTHON GO KOTLIN C# RUBY C++ Java Script

Example :

package main
import "fmt"

var root *Node

type Node struct {

    data int
    left *Node
    right *Node
    height int
    size int
}

func createNode(data int) *Node {

	var node Node

    node.data = data
    node.left = nil
    node.right = nil
    node.height = 1 
    
    return &node
}

    
func inOrder(root *Node) {
    if (root != nil) {
        inOrder(root.left);
        fmt.Print(root.data, "     ");
        inOrder(root.right)
    }
}

func leftRotate(rootNode *Node) *Node {

    newRoot := rootNode.right
    rootNode.right = rootNode.right.left
    newRoot.left = rootNode
    rootNode.height = setHeight(rootNode)
    rootNode.size = setSize(rootNode)
    newRoot.height = setHeight(newRoot)
    newRoot.size = setSize(newRoot)
    return newRoot
}

func rightRotate(rootNode *Node) *Node {

    newRoot := rootNode.left
    rootNode.left = rootNode.left.right
    newRoot.right = rootNode
    rootNode.height = setHeight(rootNode)
    rootNode.size = setSize(rootNode)
    newRoot.height = setHeight(newRoot)
    newRoot.size = setSize(newRoot)
    return newRoot
    }

func setSize(rootNode *Node) int {

	var tempValLeft int
	var tempValRight int
		
	var tempValMax int
		
    if(rootNode == nil) {
        return 0;
    }
        
    if (rootNode.left != nil) {
        tempValLeft = rootNode.left.size
    } else {
        tempValLeft = 0
    }
        
    if (rootNode.right != nil) {
        tempValRight = rootNode.right.size
    } else {
        tempValRight = 0
    }
        
    if (tempValLeft > tempValRight) {
        tempValMax = tempValLeft
    } else {
        tempValMax = tempValRight
    } 
    return 1 + tempValMax
}

func insert(rootNode *Node, data int) *Node {

    if (rootNode == nil) {
        rootNode = createNode(data)
        return rootNode
    }

    if (data < rootNode.data) {
        rootNode.left = insert(rootNode.left, data)
    } else if (data > rootNode.data) {
        rootNode.right = insert(rootNode.right, data)
    }    

    balanceFactor := balance(rootNode.left, rootNode.right)

    if(balanceFactor > 1) {

        if(height(rootNode.left.left) >= height(rootNode.left.right)){

            rootNode = rightRotate(rootNode)
        } else {

            rootNode.left = leftRotate(rootNode.left)
            rootNode = rightRotate(rootNode)
        }
    } else if(balanceFactor < -1) {

        if(height(rootNode.right.right) >= height(rootNode.right.left)) {

            rootNode = leftRotate(rootNode)
        } else {

            rootNode.right = rightRotate(rootNode.right)
            rootNode = leftRotate(rootNode)
        }
    } else {

        rootNode.height = setHeight(rootNode)
        rootNode.size = setSize(rootNode)
    }
    return rootNode
}

func balance(rootNodeLeft *Node, rootNodeRight *Node) int {
    return height(rootNodeLeft) - height(rootNodeRight)
}

func setHeight(rootNode *Node) int {
    
    var tempValLeft int
	var tempValRight int
		
	var tempValMax int
		
    if(rootNode == nil) {
        return 0;
    }
        
    if (rootNode.left != nil) {
        tempValLeft = rootNode.left.height
    } else {
        tempValLeft = 0
    }
        
    if (rootNode.right != nil) {
        tempValRight = rootNode.right.height
    } else {
        tempValRight = 0
    }
        
    if (tempValLeft > tempValRight) {
        tempValMax = tempValLeft
    } else {
        tempValMax = tempValRight
    } 
    return 1 + tempValMax
}

func height(rootNode *Node) int {

    if(rootNode == nil){
        return 0
    } else {
        return rootNode.height
    }
}

func main() {
        
    root = insert(root, 16)
    root = insert(root, 10)
    root = insert(root, 25)
    root = insert(root, 2)
    root = insert(root, 13)
    root = insert(root, 18)
    root = insert(root, 30)
    root = insert(root, 7)
    root = insert(root, 8)

    fmt.Printf("Inorder Traversal : \n")
    inOrder(root)
    
    fmt.Println()
}

Output :

Inorder Traversal :
2 7 8 10 13 16 18 25 30

So, we will constructing the below Balanced Binary Search Tree, by inserting the elements one by one.

And we have the below methods :

func insert(rootNode *Node, data int) *Node

This method is used to insert a new Node to the Binary Search Tree. Then balance it as per the rules of AVL Tree.

func leftRotate(rootNode *Node) *Node

This method is used to Left Rotate the Binary Search Tree in order to balance it as per the rules of AVL Tree.

func rightRotate(rootNode *Node) *Node

This method is used to Right Rotate the Binary Search Tree in order to balance it as per the rules of AVL Tree.

func balance(rootNodeLeft *Node, rootNodeRight *Node) int

As the name suggests, this method is used to balance the Binary Search Tree as per the rules of AVL Tree.

func setHeight(rootNode *Node) int and func height(rootNode *Node) int

As we know an AVL Tree is a height balanced Binary Search Tree. And the int setHeight(...) and int height(...) are the methods where the height of the AVL Tree is calculated.

Code explanation for AVL Tree :

Firstly, we would be constructing the above Binary Search Tree, with the help of insert method.

root = insert(root, 16)

Initially, the variable root is nil. So, we take the element i.e. 16 in this case. And pass the root and the element 16 to the insert method.

func insert(rootNode *Node, data int) *Node {

	if (rootNode == nil) {
		rootNode = createNode(data)
		return rootNode
	}

	if (data < rootNode.data) {
		rootNode.left = insert(rootNode.left, data)
	} else if (data > rootNode.data) {
		rootNode.right = insert(rootNode.right, data)
	}

	balanceFactor := balance(rootNode.left, rootNode.right)

	if(balanceFactor > 1) {

		if(height(rootNode.left.left) >= height(rootNode.left.right)){

			rootNode = rightRotate(rootNode)
		} else {

			rootNode.left = leftRotate(rootNode.left)
			rootNode = rightRotate(rootNode)
		}
	} else if(balanceFactor < -1) {

		if(height(rootNode.right.right) >= height(rootNode.right.left)) {

			rootNode = leftRotate(rootNode)
		} else {

			rootNode.right = rightRotate(rootNode.right)
			rootNode = leftRotate(rootNode)
		}
	} else {

		rootNode.height = setHeight(rootNode)
		rootNode.size = setSize(rootNode)
	}
	return rootNode
}

The first line of the insert(...) checks if the rootNode is null or not.

if (rootNode == nil) {
	rootNode = createNode(data)
	return rootNode
}

And in this case rootNode is nil. So, we create a new Node object with the data i.e. 16. And the Node object is created with the below method createNode(...).

func createNode(data int) *Node {
	var node Node
	node.data = data
	node.left = nil
	node.right = nil
	node.height = 1
	return &node
}

Then assign it to the Node object root.

And return the Node node.

return &node

In the next line we try to insert the element 13 in the Binary Search Tree using the insert(...) method.

root = insert(root, 16)

And the insert(...) method is called.

And in the similar way we check,

if (data < rootNode.data) {
	rootNode.left = insert(rootNode.left, data)
} else if (data > rootNode.data) {
	rootNode.right = insert(rootNode.right, data)
}

If the data i.e. 13 is less than the root element root.data. And the first condition matches.

if (data < rootNode.data) {
	rootNode.left = insert(rootNode.left, data)
}

And a recursive call is made to insert(...) method.

rootNode.left = insert(rootNode.left, data)

So, the method executes and sets the value of rootNode.left with this Node that contains 13.

And since we are dealing with AVL tree, it is time to calculate the Balance Factor.

balanceFactor := balance(rootNode.left, rootNode.right)

And the balance(...) method is called.

func balance(rootNodeLeft *Node, rootNodeRight *Node) int {
	height(rootNodeLeft) - height(rootNodeRight)
}

And the balance(...) method is quite simple to understand. It calls the height(...) method for the Left Node and Right Node and calculates the difference between Left and Right Node.

And as we know, Balance Factor is the difference of Left and Right Node.

And now since we got the Balance Factor, we need to perform Left or Right rotation if the value of Balance Factor is other than -1, 0 or 1.

if(balanceFactor > 1) {

	if(height(rootNode.left.left) >= height(rootNode.left.right)){

		rootNode = rightRotate(rootNode)
	} else {

		rootNode.left = leftRotate(rootNode.left)
		rootNode = rightRotate(rootNode)
	}
} else if(balanceFactor < -1) {

	if(height(rootNode.right.right) >= height(rootNode.right.left)) {

		rootNode = leftRotate(rootNode)
	} else {

		rootNode.right = rightRotate(rootNode.right)
		rootNode = leftRotate(rootNode)
	}
} else {

	rootNode.height = setHeight(rootNode)
	rootNode.size = setSize(rootNode)
}

So, if the Balance Factor is greater than 1.

if(balanceFactor > 1) {

	if(height(rootNode.left.left) >= height(rootNode.left.right)){

		rootNode = rightRotate(rootNode)
	} else {

		rootNode.left = leftRotate(rootNode.left)
		rootNode = rightRotate(rootNode)
	}
}

We go to the left child of the rootNode and try checking, if the height of the left child of the left node is greater than the height of the right child of the left node.

if(height(rootNode.left.left) >= height(rootNode.left.right)){

	rootNode = rightRotate(rootNode)
}

If the above condition satisfies, then we can assume that we need to right rotate the tree.

And the rightRotate(...) method is called to balance the tree.

func leftRotate(rootNode *Node) *Node {

	newRoot := rootNode.right
	rootNode.right = rootNode.right.left
	newRoot.left = rootNode
	rootNode.height = setHeight(rootNode)
	rootNode.size = setSize(rootNode)
	newRoot.height = setHeight(newRoot)
	newRoot.size = setSize(newRoot)
	return newRoot
}

Let us take the below example to understand the rightRotate(...) method.

Say, rootNode has the element 13.

Now, we create a Node named newNode and assign the Node rootNode.left to it.

newRoot := rootNode.right

And as we can see in the above diagram rootNode.left has the element 7 in it.

Then we take the rootNode.left.right (i.e. The right child of 7, which is null in the above case) and assign it to rootNode.left.

rootNode.left = rootNode.left.right

And the below tree looks like,

Next, we take the rootNode and assign it as the right child of the newly created Node newRoot.

newRoot.right = rootNode

So, the right rotation is performed,

The rest few lines of the method is quite easy to understand.

The height is set by calling the setHeight(...) method.

rootNode.height = setHeight(rootNode)

And finally, the Node newRoot is returned.

return newRoot

And continuing this way we get the below AVL tree.

❮❮ PREV NEXT ❯❯