Balanced Tree

When using a tree data structure, it's common for the tree to become unbalanced over time due to the insertion order of nodes, which can in turn affect the performance of our programs. Let's define a balanced tree as one where the difference in height of the left and right subtrees is at most one, for all nodes in the given tree. Write a function is_balanced(node) that determines whether a binary tree is balanced or not.

Input: The root node of a binary tree

Output: True if the tree is balanced, False otherwise.

Assume you are given the root node of a tree that conforms to the following interface:


class Node {
  left: Node
  right: Node
  value: any
}

Examples

Example 1: Balanced


Tree:
      a
     / \
    b   c
   / \   \
  d   e   f

is_balanced(a) # => True

Example 2: Balanced


Tree:
     a
    / \
   b   c
        \
         d

is_balanced(a) # => True

Example 3: Not Balanced


Tree:
     a
    / \
   b   c
        \
         d
          \
           e

is_balanced(a) # => False

Example 4: Not Balanced


Tree:
      a
     / \
    b   c
   /     \
  d       e
 /         \
f           g

is_balanced(a) # => False

Note that while the last tree seems symmetrical, it is not balanced because nodes b and c are not balanced.

We can calculate the depth (or height) of a tree by counting the number of nodes from the root to the leaf.

You might find it helpful to use recursion to traverse the nodes and find their height.

For every node, you should check whether the height of its subtrees meets the definition of being balanced.

How are you calculating the height of the tree? Does your function explore all paths?

Does your solution reject trees like the last example above?

Can you optimize the performance of your function? Can you find a linear time solution?

This is a classic data structures question that really tests your understanding of concepts like trees, recursion, and algorithmic runtime. The quality of being "balanced" is actually very important in real-world applications—that's why we created self-balancing variants such as b-trees and red-black trees.

The key to solving this problem is to read the prompt and examples carefully. Let's begin with the definition: "a balanced tree is one where the difference in height of the left and right subtrees is at most one, for all nodes in the given tree." Off the bat, we know we need to compare the height of the left and right subtrees. We also know we need to do this comparison for all nodes in the tree. Now the question is, how do we get the height?

Let's start off by writing a function to get the height of a tree. If you aren't certain what the "height" is, you could ask your interviewer to clarify. The height of a tree is typically defined as the maximum number of 'layers' from top to bottom.

def get_height(node):
    if not node:
        return 0

    l_height = get_height(node.left)
    r_height = get_height(node.right)

    return max(l_height, r_height) + 1

Now, let's use this to implement our main function, is_balanced. All we have to do is explore each node and make sure the height of its subtrees meet the condition in the prompt: the difference in height should be less than or equal to 1.

def is_balanced(node):
    if not node:
        return True

    l_height = get_height(node.left)
    r_height = get_height(node.right)
    height_diff = abs(l_height - r_height)

    if height_diff > 1:
        return False

    return is_balanced(node.left) and is_balanced(node.right)

A note on recursion: Whenever you implement a recursive function, in the back of your mind you should ask whether recursion is the best way to accomplish your goal. Recursion has the downside of using the callstack, which can lead to stack overflow for deeply nested data. You're unlikely to hit this error during an interview, but it's worth mentioning as a real-life constraint and explaining how to implement a non-recursive version (e.g. by using a stack).

Time Complexity: $O(n^2)$ . For each node, the get_height function is called to calculate the heights of the left and right subtrees. Since get_height is called for each of the n nodes in the worst case $O(n)$ , this results in $O(n * n) = O(n^2)$ time complexity.

Space Complexity: $O(h)$ . the space complexity is determined by the maximum depth of the recursion stack, which is the height of the tree ( $h$ ). Each recursive call of is_balanced adds to the call stack.

Optimized Approach

We can optimize the is_balanced function to reduce the time complexity to $O(n)$ by combining the height calculation and balance checking into a single traversal. Here’s how:

def check_height_and_balance(node):
    if not node:
        return 0, True

    l_height, l_balanced = check_height_and_balance(node.left)
    r_height, r_balanced = check_height_and_balance(node.right)

    if not l_balanced or not r_balanced:
        return max(l_height, r_height) + 1, False

    height_diff = abs(l_height - r_height)
    if height_diff > 1:
        return max(l_height, r_height) + 1, False

    return max(l_height, r_height) + 1, True

def is_balanced(node):
    _, balanced = check_height_and_balance(node)
    return balanced

Can you use a stack to solve this?

Community solutions

💬 Community guidelines

Stay on topic. Use this section for submitting solutions and providing feedback to others.
Be inclusive. Exponent is a diverse community. Please respect others' opinions and beliefs.

How would you answer? Submit your response below and get feedback from the community.

Tiago R.148

November 4, 2024

JavaScript

function visitChildren(node) {
    let leftSubtreeHeight = 0;
    let rightSubtreeHeight = 0;
    let isChildrenBalanced = true;

    if (node.left) {
        const { isBalanced, height } = visitChildren(node.left);
        isChildrenBalanced = isChildrenBalanced && isBalanced;
        leftSubtreeHeight += height + 1;
    }
    if (isChildrenBalanced && node.right) {
        const { isBalanced, height } = visitChildren(node.right);
        isChildrenBalanced = isChildrenBalanced && isBalanced;
        rightSubtreeHeight += height + 1;
    }

    return {
        isBalanced: isChildrenBalanced && (Math.abs(leftSubtreeHeight - rightSubtreeHeight) < 2),
        height: Math.max(leftSubtreeHeight, rightSubtreeHeight),
    };
}

function isBalanced(node) {
    if (node) {
        return visitChildren(node).isBalanced;
    }
    return false;
}

Basil A.351

January 27, 2025

Java

public class Solution {

    // Definition for a binary tree node
    static class Node {
        String value;
        Node left;
        Node right;

        Node(String value) {
            this.value = value;
            this.left = null;
            this.right = null;
        }

        Node(String value, Node left, Node right) {
            this.value = value;
            this.left = left;
            this.right = right;
        }
    }

    static boolean balanced;

    public static boolean isBalanced(Node node) {
        // your code goes here

        if(node == null) {
            return true;
        }

        balanced = true;

        height(node);


        return balanced;
    }

    static int height(Node node) {
        if(node == null) {
            return 0;
        }

        int l = height(node.left);
        int r = height(node.right);

        
        if(Math.abs(l-r) > 1) {
            balanced = false;
        }

        return 1 + Math.max(l, r);
    }

    public static void main(String[] args) {
        // debug your code below
        Node a = new Node("a", 
                          new Node("b", 
                                   new Node("d"), 
                                   new Node("e")),
                          new Node("c", 
                                   null, 
                                   new Node("f")));

        System.out.println(isBalanced(a)); // Should return true for a balanced tree
    }
}

Sahil A.

Member

February 12, 2025

C++

#include <iostream>
#include <algorithm>
#include <queue>
    
class Node {
public:
    char value;
    Node *left;
    Node *right;

    Node(char val, Node *leftChild = nullptr, Node *rightChild = nullptr)
        : value(val), left(leftChild), right(rightChild) {}
};

bool is_balanced(Node *node) {
    // your code goes here
    if(!node){
        return true;
    }

    if(!(node->left)&&!(node->right)){
        return true;
    }

    if(!(node->left)&&(node->right!=nullptr)){
        if(node->right->left!=nullptr||node->right->right!=nullptr){
            return false;
        }
        else{
            return is_balanced(node->right);
        }
    }
    else if(!(node->right)&&(node->left!=nullptr)){
        if(node->left->left!=nullptr||node->left->right!=nullptr){
            return false;
        }
        else{
            return is_balanced(node->left);
        }
    }
    else{
        return is_balanced(node->left) && is_balanced(node->right);
    }

}

// comment out main() before running tests
int main() {
    // debug your code below
    Node *root = new Node('a');
    root->left = new Node('b');
    root->right = new Node('c');
    root->left->left = new Node('d');
    root->left->right = new Node('e');
    root->right->right = new Node('f');

    // Test is_balanced function
    std::cout << is_balanced(root) << std::endl; // Output: 1 (true)

    // Free dynamically allocated memory
    delete root->left->left;
    delete root->left->right;
    delete root->right->right;
    delete root->left;
    delete root->right;
    delete root;

    return 0;
}

Dana L.75

February 13, 2025

Python

class Node:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def get_depths(node):
    nodes = []
    if node.left:
        depths = get_depths(node.left)
        for depth in depths:
            nodes.append(1 + depth)
    else:
        nodes.append(0)
        
    if node.right:
        depths = get_depths(node.right)
        for depth in depths:
            nodes.append(1 + depth)
    else:
        nodes.append(0)

    return nodes

def is_balanced(node):
    if node is None:
        return True
    depths = get_depths(node)
    min_d = min(depths)
    max_d = max(depths)
    if min_d == max_d or min_d + 1 == max_d:
        return True
    else:
        return False
    
# debug your code below
root = Node(1)
root.left = Node(2)
root.right = Node(3)
root.left.left = Node(4)
root.left.right = Node(5)
root.right.right = Node(6)
print(is_balanced(root))

See community answers

ML Engineer Introduction

ML System Design Questions

Overview

Mock Interviews & Practice Questions

ML Concepts Questions

Overview

Model & Algorithm Fundamentals

Mock Interviews

ML Coding Questions

Overview

Mock Interviews & Practice Questions

Software Engineering Coding Questions

Overview

Fast Track

Time & Space Complexity

Coding Patterns

Arrays

Hash Tables

Searching & Sorting

String Manipulation

Graphs

Trees

Stacks & Queues

Linked Lists

Heaps

Recursion & Backtracking

Dynamic Programming

Behavioral Interviews: A Practical Guide for Engineers

Fast Track (Start Here)

Theory

Tactics

Mock Interviews & Analyses

Practice

ML Engineer Introduction

ML System Design Questions

Overview

Mock Interviews & Practice Questions

ML Concepts Questions

Overview

Model & Algorithm Fundamentals

Mock Interviews

ML Coding Questions

Overview

Mock Interviews & Practice Questions

Software Engineering Coding Questions

Overview

Fast Track

Time & Space Complexity

Coding Patterns

Arrays

Hash Tables

Searching & Sorting

String Manipulation

Graphs

Trees

Stacks & Queues

Linked Lists

Heaps

Recursion & Backtracking

Dynamic Programming

Behavioral Interviews: A Practical Guide for Engineers

Fast Track (Start Here)

Theory

Tactics

Mock Interviews & Analyses

Practice

Examples

Optimized Approach