CheckIfBinaryTreeBalanced.java

package com.thealgorithms.datastructures.trees;

import java.util.HashMap;
import java.util.Stack;

/**
 * This class will check if a BinaryTree is balanced. A balanced binary tree is
 * defined as a binary tree where the difference in height between the left and
 * right subtree of each node differs by at most one.
 * <p>
 * This can be done in both an iterative and recursive fashion. Below,
 * `isBalancedRecursive()` is implemented in a recursive fashion, and
 * `isBalancedIterative()` is implemented in an iterative fashion.
 *
 * @author [Ian Cowan](<a href="https://github.com/iccowan">Git-Ian Cowan</a>)
 */
public final class CheckIfBinaryTreeBalanced {
 private CheckIfBinaryTreeBalanced() {
 }
 /**
 * Recursive is BT balanced implementation
 *
 * @param root The binary tree to check if balanced
 */
 public static boolean isBalancedRecursive(BinaryTree.Node root) {
 if (root == null) {
 return true;
 }
 // Create an array of length 1 to keep track of our balance
 // Default to true. We use an array, so we have an efficient mutable object
 boolean[] isBalanced = new boolean[1];
 isBalanced[0] = true;

 // Check for balance and return whether we are balanced
 isBalancedRecursive(root, 0, isBalanced);
 return isBalanced[0];
 }

 /**
 * Private helper method to keep track of the depth and balance during
 * recursion. We effectively perform a modified post-order traversal where
 * we are looking at the heights of both children of each node in the tree
 *
 * @param node The current node to explore
 * @param depth The current depth of the node
 * @param isBalanced The array of length 1 keeping track of our balance
 */
 private static int isBalancedRecursive(BinaryTree.Node node, int depth, boolean[] isBalanced) {
 // If the node is null, we should not explore it and the height is 0
 // If the tree is already not balanced, might as well stop because we
 // can't make it balanced now!
 if (node == null || !isBalanced[0]) {
 return 0;
 }

 // Visit the left and right children, incrementing their depths by 1
 int leftHeight = isBalancedRecursive(node.left, depth + 1, isBalanced);
 int rightHeight = isBalancedRecursive(node.right, depth + 1, isBalanced);

 // If the height of either of the left or right subtrees differ by more
 // than 1, we cannot be balanced
 if (Math.abs(leftHeight - rightHeight) > 1) {
 isBalanced[0] = false;
 }

 // The height of our tree is the maximum of the heights of the left
 // and right subtrees plus one
 return Math.max(leftHeight, rightHeight) + 1;
 }

 /**
 * Iterative is BT balanced implementation
 */
 public static boolean isBalancedIterative(BinaryTree.Node root) {
 if (root == null) {
 return true;
 }

 // Default that we are balanced and our algo will prove it wrong
 boolean isBalanced = true;

 // Create a stack for our post order traversal
 Stack<BinaryTree.Node> nodeStack = new Stack<>();

 // For post order traversal, we'll have to keep track of where we
 // visited last
 BinaryTree.Node lastVisited = null;

 // Create a HashMap to keep track of the subtree heights for each node
 HashMap<BinaryTree.Node, Integer> subtreeHeights = new HashMap<>();

 // We begin at the root of the tree
 BinaryTree.Node node = root;

 // We loop while:
 // - the node stack is empty and the node we explore is null
 // AND
 // - the tree is still balanced
 while (!(nodeStack.isEmpty() && node == null) && isBalanced) {
 // If the node is not null, we push it to the stack and continue
 // to the left
 if (node != null) {
 nodeStack.push(node);
 node = node.left;
 // Once we hit a node that is null, we are as deep as we can go
 // to the left
 } else {
 // Find the last node we put on the stack
 node = nodeStack.peek();

 // If the right child of the node has either been visited or
 // is null, we visit this node
 if (node.right == null || node.right == lastVisited) {
 // We assume the left and right heights are 0
 int leftHeight = 0;
 int rightHeight = 0;

 // If the right and left children are not null, we must
 // have already explored them and have a height
 // for them so let's get that
 if (node.left != null) {
 leftHeight = subtreeHeights.get(node.left);
 }

 if (node.right != null) {
 rightHeight = subtreeHeights.get(node.right);
 }

 // If the difference in the height of the right subtree
 // and left subtree differs by more than 1, we cannot be
 // balanced
 if (Math.abs(rightHeight - leftHeight) > 1) {
 isBalanced = false;
 }

 // The height of the subtree containing this node is the
 // max of the left and right subtree heights plus 1
 subtreeHeights.put(node, Math.max(rightHeight, leftHeight) + 1);

 // We've now visited this node, so we pop it from the stack
 nodeStack.pop();
 lastVisited = node;

 // Current visiting node is now null
 node = null;
 // If the right child node of this node has not been visited
 // and is not null, we need to get that child node on the stack
 } else {
 node = node.right;
 }
 }
 }

 // Return whether the tree is balanced
 return isBalanced;
 }
}