Implement a function to check if a tree is balanced Yet another interview question:


  Implement a function to check if a tree is balanced. For the purposes of this question, a balanced tree is defined to be a tree such that no two leaf nodes differ in distance from the root by more than one.


First off, did I understand the question correctly? In other words, we are looking for a subtree in a tree (or rotation of it) that looks like:


     root
    /    \
 leaf_1  leaf_2
  /
leaf1_1



My initial idea was to count the depth of the tree and check if maxDepth - minDepth &lt; 2 (which requires 2 traversal) . But I think I found a more efficient solution. The trade-off is that we should keep the ref to the root. I assume n-ary tree.

Is this correct?

import java.util.LinkedList;

public class Trees {

    Node threeRoot;
    private boolean isBalanced = true;

    Trees(Node threeRoot){
        this.threeRoot = threeRoot;
    }

    static class  Node{
        public Node root;
        public LinkedList&lt;Node&gt; children = new LinkedList&lt;&gt;();        
    }

     boolean isBalanced(Node node){
       if (threeRoot == null) {
            isBalanced = true;
        }
        if (threeRoot.children == null || threeRoot.children.isEmpty()) {
            isBalanced = true;
        }

        isBalancedRed(node);          
        return isBalanced;
    }

     void isBalancedRed(Node node) {
        if (node.children.isEmpty()) {
            if (node.root != null &amp;&amp; node.root.children.size() == 1) {
                isBalanced = false;
            }
        }
        for(Node nodeEl : node.children)
        {
            isBalancedRed(nodeEl);
        }
    }

    public static void main(String[] args){
        Trees tree = new Trees(new Node());
        Node root = tree.threeRoot;
        Node leaf1 = new Node();
        Node leaf2 = new Node();
        Node leaf11 = new Node();

        root.children.add(leaf1);
        leaf1.root = root;
        root.children.add(leaf2);
        leaf2.root = root;

        leaf2.children.add(leaf11);
        leaf11.root = leaf2;

        System.out.println(tree.isBalanced(root));           
    }        
}
First off, I would use something other than a LinkedList to hold the children of an n-ary tree. A LinkedList requires you to enumerate the list in order to get the reference to any child other than the leftmost, which is going to slow you down.

Second, your code is basically returning false in any situation where the Node's parent has one child that itself doesn't have any children. In your example tree, that's the exact case (you have two children of the root, one of which has another child) but the tree is balanced according to the rules, so I would expect this algorithm to return a lot of false positives (and negatives).

Finding the depth of an n-ary tree is a linear-bound operation; you traverse each branch of the tree recursively, finding the depth at each leaf node. You can find both the maximum and minimum depth in one traversal; at each level, ask each branch for its maximum and minimum depth, passing the depth of the current node. The base case is that of a leaf node (no children); the min and max is the depth of that leaf. You can store this result in a Tuple or in a more specialized MinMax struct, as you please. With those results returned to the calling level, scan them to find the lowest Min and the greatest Max, and return that to the caller. 

Here's a basic implementation:

public class MinMax&lt;T&gt;
{
    private T minVal;
    private T maxVal;

    public MinMax(T min, T max)
    {
       minVal = min;
       maxVal = max;
    }

    public T min() {return minVal;}
    public T max() {return maxVal;}
}

private MinMax&lt;Integer&gt; GetMinMaxDepth(Node node, int currDepth)
{
    //base case; no children. Max and min depth at this leaf is currDepth.
    if(node.children == null || node.children.size() == 0)
       return new MinMax&lt;Integer&gt;(currDepth, currDepth);

    //determine the maximum and minimum depth of all child branches,
    //and track the absolute max and min across all of them.
    int min = Integer.MAX_VALUE, max = -1;
    for(Node nextNode : node.children)
    {
        MinMax&lt;Integer&gt; childDepth = GetMinMaxDepth(nextNode, currDepth+1);
        if(childDepth.min() &lt; min) min = childDepth.min();
        if(childDepth.max() &gt; max) max = childDepth.max();
    }      

    return new MinMax&lt;Integer&gt;(min, max);
}

boolean isBalanced()
{
    //easy cases; a tree with 0 or 1 elements.
    if(treeRoot == null || treeRoot.children == null || treeRoot.children.size() == 0)
       return true;
    //traverse the tree and find the minimum and maximum depth.
    MinMax&lt;Integer&gt; minMaxDepth = GetMinMaxDepth(treeRoot, 0);

    return minMaxDepth.max() - minMaxDepth.min() &gt; 1;
}


This could probably be optimized in our case to return early if we discover that any node's child depths differ by more than 1. We could hack it by throwing an Exception, but the proper strategy is almost as easy to implement:

private MinMax&lt;Integer&gt; GetMinMaxDepth(Node node, int currDepth)
{
    //base case; no children. max and min depth at this leaf is currDepth.
    if(node.children == null || node.children.size() == 0)
       return new MinMax&lt;Integer&gt;(currDepth, currDepth);

    //determine the maximum and minimum depth of all child branches,
    //and track the absolute max and absolute minimum.
    int min = Integer.MIN_VALUE, max = -1;
    for(Node nextNode : node.children)
    {
        MinMax&lt;Integer&gt; childDepth = GetMinMaxDepth(nextNode, currDepth+1);
        //Exit now if child is unbalanced
        if(childDepth.max() - childDepth.min() &gt; 1) return childDepth;
        if(childDepth.min() &lt; min) min = childDepth.min();
        if(childDepth.max() &gt; max) max = childDepth.max();
        //Exit now if current node is unbalanced
        if(max-min &gt; 1) break;
    }      

    return new MinMax&lt;Integer&gt;(min, max);
}


The upside is that we quit as soon as we know the answer to the question (is the tree unbalanced?), which will increase the average performance (but not the worst-case performance on a tree that is balanced or is unbalanced at its furthest extremities). The downside is that we no longer know how unbalanced the tree is; the MaxMin returned to the top level will always have a Max and Min that differ by the first detected difference greater than the threshold (probably 2), not the absolute difference in depth of leaf nodes in the tree.

The one case that is difficult to determine in an n-ary tree is that of a tree that never forks. This algorithm will find a depth difference between any two or more branches, but when there is only one branch, the maximum and minimum depths of the tree are the same and there's nothing to use for comparison. Technically, it would have only one "leaf" node (usually defined as a node with no children) and so by the given definition all (one) of the leaves are the same distance from the root, however if you looked at the map of an N-ary tree that was more like a linked list (or even a V or Y shape) and went deeper than two levels, you wouldn't call it balanced.

Perhaps a change in definition may be required; a node is a "leaf" node if it does not have N child nodes (where N is the order of the N-ary tree). So, in a quaternary tree (4 children per node), if any node has fewer than 4 children, then the depth of the current node is considered as a "minimum" depth and will almost certainly be the minimum depth at that level. That's another easy change:

private MinMax&lt;Integer&gt; GetMinMaxDepth(Node node, int currDepth)
{
    //base case; no children. max and min depth at this leaf is currDepth.
    if(node.children == null || node.children.size() == 0)
       return new MinMax&lt;Integer&gt;(currDepth, currDepth);

    //determine the maximum and minimum depth of all child branches,
    //and track the absolute max and absolute minimum.
    //Exit as soon as we know that this or any child node is unbalanced.
    int min = Integer.MIN_VALUE, max = -1;

    //If this node isn't "full", its depth is the minimum depth for this branch
    if(node.children.size &lt; max_children) min = currDepth;      

    for(Node nextNode : node.children)
    {
        MinMax&lt;Integer&gt; childDepth = GetMinMaxDepth(nextNode, currDepth+1);
        //check if child is already unbalanced
        if(childDepth.max() - childDepth.min() &gt; 1) return childDepth;
        if(childDepth.min() &lt; min) min = childDepth.min();
        if(childDepth.max() &gt; max) max = childDepth.max();
        //check if current node is unbalanced
        if(max-min &gt; 1) break;
    }      

    return new MinMax&lt;Integer&gt;(min, max);
}