Balance a Binary Search Tree

Good evening. It is Tuesday January 12, 2021. We are still in the middle of the COVID-19 pandemic. For some reason vaccination in the USA and specific in Minnesota is not progressing at a reasonable pace. Hopefully when all political issues clear, things will start progressing at a better pace.

The main subject for this post is LeetCode 1382 Balance a Binary Search Tree. A binary search tree is a regular binary tree with an additional condition. Given any node in the tree, the values on the left sub tree will be smaller that the value of the node and all values in the right sub tree will be larger than the value in the node. This holds true for all nodes if they have a left and / or right sub trees.

In general if you insert nodes in a tree with random values the resulting BST should be almost balanced. If you feed values in ascending or descending order, you end up with a lopped tree in which one end will resemble a queue. Access time will be in O(n) time instead of O(log(n)).

Given a binary search tree, return a balanced binary search tree with the same node values.

A binary search tree is balanced if and only if the depth of the two subtrees of every node never differ by more than 1.

If there is more than one answer, return any of them.

Constraints:

o The number of nodes in the tree is between 1 and 10^4.
o The tree nodes will have distinct values between 1 and 10^5.

In this problem, we are given a BST and are required to generate a new BST as balanced as possible. Note that the operation would be invoked once. We do not have to implement an AVL tree or a Red-Black tree algorithm. Both of these algorithms attempt to balance the BST after a node insertion or deletion.

To be honest with you I have not implemented a one of the mentioned algorithms in a while. So I decided to use the AVL algorithm as a first pass. At some point I completed the code but decided not to submit it because I needed an additional member in the TreeNode.

On a second implementation I used a different approach and it was submitted and accepted. You will see the changes I did to improve on its performance.

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    public TreeNode balanceBST(TreeNode root) {
        
    }
}

LeetCode provided the TreeNode class and the signature for the balanceBST() function / method. We are given a BST and we need to returned a balanced BST with the same node values as the one given as an argument.

Note that there is no height member in the TreeNode class.

In order to refresh the AVL algorithm I consulted three text books that I have in my physical library. The explanation was not optimal. I recall that while attending college I had a light blue text book with very good descriptions of different algorithms including AVL trees.

I also watched the associated section in Advanced Data Structures and Algorithms in Java 9 by Debasish Ray Chawdhuri in Packt published by O’Reilly. The descriptions are quite good. The code is somewhat convoluted.

We are going to solve the problem using Java in a Windows 10 computer using the VSCode IDE. This will require for us to build some kind of test scaffolding. When done we can copy and paste the code for the solution into the LeetCode IDE and submit it for approval.

1,null,2,null,3,null,4,null,null
main <<< arr:[1, null, 2, null, 3, null, 4, null, null]
main <<< bst levelOrder:
1
2
3
4
main <<< avl levelOrder:
2
1,3
4
main <<< avl inOrder: (1,1) (2,3) (3,2) (4,1)
main <<< balanced levelOrder:
2
1,3
4
main <<< balanced inOrder: (1,1) (2,1) (3,1) (4,1)


5,4,null,3,null,2,null,1,null
main <<< arr:[5, 4, null, 3, null, 2, null, 1, null]
main <<< bst levelOrder:
5
4
3
2
1
main <<< avl levelOrder:
2
1,4
3,5
main <<< avl inOrder: (1,1) (2,3) (3,1) (4,2) (5,1)
main <<< balanced levelOrder:
3
1,4
2,5
main <<< balanced inOrder: (1,1) (2,1) (3,1) (4,1) (5,1)


12,8,18,5,11,17,null,4,7,null,null,null,null,2,null,null,null
main <<< arr:[12, 8, 18, 5, 11, 17, null, 4, 7, null, null, null, null, 2, null, null, null]       
main <<< bst levelOrder:
12
8,18
5,11,17
4,7
2
main <<< avl levelOrder:
7
4,11
2,5,8,17
12,18
main <<< avl inOrder: (2,1) (4,2) (5,1) (7,4) (8,1) (11,3) (12,1) (17,2) (18,1) 
main <<< balanced levelOrder:
8
4,12
2,5,11,17
7,18
main <<< balanced inOrder: (2,1) (4,1) (5,1) (7,1) (8,1) (11,1) (12,1) (17,1) (18,1)

The LeetCode site provides a single test case. It is the first case in the screen capture. The other two I generated them and edited based on what I wanted to test. At some point I had over 10 different tests. I decided to condense them into three.

It seems that we are provided the valued for the nodes of our BST in depth first order traversal. Our test scaffolding seems to read the input line and generate an arr[] with the specified values.

It seems that we populate the BST and display it in level order. This is a typical example of a BST with all the nodes on the left sub tree. The search performance turns to be O(n) since we are traversing what is equivalent to a linked list.

It appears that we use our BST and generate an AVL BST. We display the AVL tree in level order. We then display the same AVL BST but in order. Note that the values are in ascending order as expected. We also display the height value which I believe we cannot include in any LeetCode solutions. In other words the first attempt is just to experiment with AVL trees which is NOT PART OF THE SOLUTION.

The next line displays a balanced BST. This tree is generated with a solution that we can submit to LeetCode. Note in the next line that the height member is just set to 1 and not used in the implementation of the solution. The reason we see a 1 and not a 0 will become apparent when we look at the code for the AVL approach.

    /**
     * Test scaffolding
     * 
     * @throws IOException
     */
    public static void main(String[] args) throws IOException {
        
        // **** open buffered reader ****
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        
        // **** read input line and generate String[] ****
        String[] strArr = br.readLine().trim().split(",");

        // **** close buffered reader ****
        br.close();

        // **** generate Integer[] array from the strArr[] ****
        Integer[] arr = new Integer[strArr.length];
        for (int i = 0; i < strArr.length; i++) {
            if (strArr[i].equals("null"))
                arr[i] = null;
            else
                arr[i] = Integer.parseInt(strArr[i]);
        }

        // ???? ????
        System.out.println("main <<< arr:" + Arrays.toString(arr));

        // **** create and populate a BST ****
        BST bst = new BST();
        bst.root = bst.populate(arr);

        // ???? ????
        System.out.println("main <<< bst levelOrder:");
        System.out.println(bst.levelOrder());


        // **** 1) balance this BST ****
        BST avl = new BST();
        avl.root = balanceBST(bst.root);

        // **** display the avl BST ****
        System.out.println("main <<< avl levelOrder:");
        System.out.println(avl.levelOrder());

        // **** in order traversal of avl BST ****
        System.out.print("main <<< avl inOrder: ");
        avl.inOrder(avl.root);
        System.out.println();


        // **** 2) balance this BST ****
        BST balanced = new BST();
        balanced.root = balanceBST2(bst.root);

        // **** display level order traversal for balanced tree  ****
        System.out.println("main <<< balanced levelOrder:");
        System.out.println(balanced.levelOrder());

        // **** in order balanced tree ****
        System.out.print("main <<< balanced inOrder: ");
        balanced.inOrder(balanced.root);
        System.out.println();
    }

We open a buffered reader, read the input line into a String[] and close the buffered reader. We then generate an Integer[] which will be used to populate the BST. We then display the BST tree in level order.

The next set of steps take the BST and generate a new tree named avl. We will shortly see the approach in some detail. The avl tree is then displayed via an in order traversal. This is not part of the solution.

We then use the BST and generate a balanced tree. Be perform a level order traversal and then an in order.

    /**
     * Given a binary search tree (BST), return a balanced BST 
     * with the same node values.
     * 
     * Traverse nodes in inorder and one by one insert into the AVL tree.
     * Time complexity of this solution is O(n Log n).
     */
    static TreeNode balanceBST(TreeNode root) {

        // **** sanity check(s) ****
        if (root == null)
            return null;

        if (root.left == null && root.right == null)
            return root;

        // **** initialization ****
        avl = new BST();

        // **** traverse inorder BST inserting nodes in AVL tree ****
        copyToAVL(root);

        // **** returned balanced BST ****
        return avl.root;
    }

The balanceBST() function starts by performing some sanity checks. I believe the first one is not needed. We then create an empty binary tree. Then we make a call to copyToAVL() function. That function traverses the BST and inserts the node values into a new AVL tree. The AVL tree is then returned. More complex that needed but it should work.

    /**
     * Inorder BST traversal.
     * Recursive call.
     */
    static void copyToAVL(TreeNode node) {
        if (node != null) {

            // **** traverse left tree ****
            copyToAVL(node.left);

            // **** insert node into AVL tree ****
            avl.root = avl.insertAVL(avl.root, node.val);

            // **** traverse right tree ****
            copyToAVL(node.right);
        }
    }

The copyToAVL() function / method performs an in order traversal of the BST. It inserts a node with the proper value into an avl tree. So far I believe it is reasonable.

    /**
     * Insert node into AVL BST.
     * Recursive call.
     */
    public TreeNode insertAVL(TreeNode node, int val) {

        // **** 1) normal BST value insertion ****
        if (node == null)
            return (new TreeNode(val));

        // **** look where to insert the new node (recursion) ****
        if (val < node.val)
            node.left = insertAVL(node.left, val);
        else if (val > node.val)
            node.right = insertAVL(node.right, val);
        else 
            return node;
        
        // **** 2) update the height of the root node (use instead private local method) ****
        node.height = 1 + max(heightAVL(node.left), heightAVL(node.right));

        // **** 3) get the balance factor ****
        int balance = getBalanceAVL(node);

        // ???? ????
        // System.out.println("insertAVL <<< balance: " + balance + " val: " + val);

        // **** 4a) Left Left case ****
        if (balance > 1 && val < node.left.val)
            return rightRotate(node);

        // **** 4b) Right Right case ****
        if (balance < -1 && val > node.right.val)
            return leftRotate(node);

        // **** 4c) Left Right case ****
        if (balance > 1 && val > node.left.val) {
            node.left = leftRotate(node.left);
            return rightRotate(node);
        }

        // **** 4d) Right Left case ****
        if (balance < -1 && val < node.right.val) {
            node.right = rightRotate(node.right);
            return leftRotate(node);
        }

        // **** return node pointer ****
        return node;
    }

The insertAVL() function / method inserts a value into an AVL tree. Note the first few lines. They follow the insertion of nodes in a regular BST.

The second step is to update the height of the member in the TreeNode class object. Because we are altering the TreeNode class I believe we cannot submit this solution.

The idea for the height is to determine if the BST is or is not balanced and act accordingly.

We then get the balance. Based on it and the value of a child node, we will perform one or two rebalance operations.

The four cases follow. I would suggest for you to get a book or look it up on-line how the rebalancing operations work. It is good to see some drawings. I just do not have the time to generate them here.

When all is wet and done, we return the balanced BST.

    /**
     * Utility function to get the height of a node in an AVL BST.
     */
    private int heightAVL(TreeNode node) {
        if (node == null) {
            return 0;
        } else {
            return node.height;
        }
    }

This function is used to get the height value from a node.

    /**
     * Utility function that returns the max of two integers.
     */
    private int max(int a, int b) {
        return (a > b) ? a : b;
    }

This function is a simpler implementation for the Math.max() function in the Java Math library. This approach should run faster.

   /**
     * Utility function that returns the balance factor 
     * of the specified node in an AVL BST.
     */
    public int getBalanceAVL(TreeNode node) {
        if (node == null) {
            return 0;
        } else {
            return heightAVL(node.left) - heightAVL(node.right);
        }
    }

This function is used to compute the balance for the specified node.

    /**
     * Rotate right in BST.
     * A right rotation decreases the depth of the left subtree by 1, 
     * and increases that of the right subtree by 1.
     */
    public TreeNode rightRotate(TreeNode y) {

        // ???? ????
        System.out.println("rightRotate <<< y: " + y.toString());

        // **** initialization ****
        TreeNode x  = y.left;
        TreeNode t2 = x.right;

        // **** perform rotation ****
        x.right = y;
        y.left  = t2;

        // **** update heights ****
        y.height = max(heightAVL(y.left), heightAVL(y.right)) + 1; 
        x.height = max(heightAVL(x.left), heightAVL(x.right)) + 1; 

        // **** return new root ****
        return x;
    }

The AVL algorithm can be summarized by the rightRotate() and leftRotate() functions / methods.

    /**
     * Rotate left in BST.
     * A left rotation decreases the depth of the right subtree by 1, 
     * and increases that of the left subtree by 1. 
     */
    public TreeNode leftRotate(TreeNode x) {

        // **** initialization ****
        TreeNode y  = x.right;
        TreeNode t2 = y.left;

        // **** perform rotation ****
        y.left  = x;
        x.right = t2;

        // **** update heights ****
        x.height = max(heightAVL(x.left), heightAVL(x.right)) + 1; 
        y.height = max(heightAVL(y.left), heightAVL(y.right)) + 1; 
  
        // **** return new root ****
        return y;
    }

This is the rightRotate() function / method.

/**
 * Enhanced TreeNode class typically used
 * to solve LeetCode binary tree problems.
 */
public class TreeNode {
    
    // **** class members ****
    int         val;
    int         height;

    TreeNode    left;
    TreeNode    right;

    /**
     * Constructor.
     */
    public TreeNode() {
    }

    /**
     * Constructor.
     */
    public TreeNode(int val) {
        this.val    = val;
        this.height = 1;
    }

    /**
     * Constructor.
     */
    public TreeNode(int val, TreeNode left, TreeNode right) {
        this.val    = val;
        this.left   = left;
        this.right  = right;
    }

    /**
     * Return a string representation of this node.
     */
    @Override
    public String toString() {
        return "(" + val + "," + height + ")";
    }
}

As you can see, we had to add the height field to the TreeNode class. That variable is only used by the AVL associated functions which we will not use while developing the actual solution. For simplicity I just left it in the class. Since the TreeNode class was not altered in LeetCode, the solution that follows works fine.

    /**
     * Given a binary search tree (BST), return a balanced BST 
     * with the same node values.
     * 
     * Runtime: 2 ms, faster than 99.69% of Java online submissions.
     * Memory Usage: 41.2 MB, less than 88.60% of Java online submissions.
     *
     * !!!! PART OF SOLUTION !!!!
     */
    static TreeNode balanceBST2(TreeNode root) {

        // **** initialization ****
        List<TreeNode> al = new ArrayList<>();

        // **** populate array list by performing an in order traversal of the BST O(n) ****
        fillALInOrder(root, al);

        // ???? ????
        // System.out.println("balanceBST2 <<<  al: " + al.toString());

        // // **** declare and fill array with results of in order traversal O(n) ****
        // int[] arr = al.stream().mapToInt(Integer::valueOf).toArray();

        // ???? ????
        // System.out.println("balanceBST2 <<< arr: " + Arrays.toString(arr));

        // **** populate and return balanced BST O(n) ****
        return populateBST(al, 0, al.size() - 1);
    }

We will use an approach that is similar to a binary search. We will populate and array list with the node values of the BST. This is done by the fillALInOrder() function / method.

Note that a few lines have been commented out. In my first passes I was going to convert the array list to an int[] and then use it in the populateBST() method. It turns out that the code runs faster when we use the array list than converting it to an int[] and then using it instead.

  /**
     * Fill the array list with the ordered node values.
     * Recursive method.
     *
     * !!!! PART OF SOLUTION !!!!
     */
    static void fillALInOrder(TreeNode root, List<TreeNode> al) {

        // **** base condition ****
        if (root == null)
            return;

        // **** visit left tree ****
        fillALInOrder(root.left, al);

        // **** insert value into array list ****
        al.add(root);

        // **** visit right tree ****
        fillALInOrder(root.right, al);
    }

This function / method is used to populate the array list with the node values in the BST while performing an in order tree traversal.

    /**
     * Populate balanced BST using int[] arr
     * Recursive call. O(n)
     *
     * !!!! PART OF SOLUTION !!!!
     */
    static TreeNode populateBST(List<TreeNode> al, int left, int right) {

        // **** base condition ****
        if (left > right)
            return null;

        // **** initialization ****
        int mid = (left + right) / 2;

        // **** ****
        TreeNode root = al.get(mid);

        // **** process left side ****
        root.left = populateBST(al, left, mid - 1);

        // **** process right side ****
        root.right = populateBST(al, mid + 1, right);

        // **** return root of balanced BST ****
        return root;
    }

This is the core of the algorithm. We use a sub section of the array list on each recursive pass. We recursively call the same method with a left and a right subset.

When all is set and done, we have created a very well balanced BST. Not only that but we did not have the need to use the height field in the TreeNode class.

Note that this is not a substitute for ateh AVL or RED-Black tree algorithms. Those are designed to constantly (when needed) balance the existing BST. In this case we balanced the BST at a greater cost.

There is additional code I used to load the data and run experiments. All that code is in my GitHub repository if you are interested.

Hope you enjoyed solving this problem as much as I did. The entire code for this project can be found in my GitHub repository.

If you have comments or questions regarding this, or any other post in this blog, or if you would like for me to help out with any phase in the SDLC (Software Development Life Cycle) of a project associated with a product or service, please do not hesitate and leave me a note below. If you prefer, send me a private e-mail message. I will reply as soon as possible.

Keep on reading and experimenting. It is the best way to learn, become proficient, refresh your knowledge and enhance your developer toolset.

One last thing, many thanks to all 5,750 subscribers to this blog!!!

Keep safe during the COVID-19 pandemic and help restart the world economy. I believe we can all see the light at the end of the tunnel.

Regards;

John

john.canessa@gmail.com

Leave a Reply

Your email address will not be published. Required fields are marked *

This site uses Akismet to reduce spam. Learn how your comment data is processed.