parikshan

balanced-binary-tree

Difficulty: Easy  ·  Topic: Trees  ·  Practice it on parikshan: /practice/balanced-binary-tree/start

The story

A trading firm runs an internal observability tool that watches their order-matching B-tree indexes in real time. When inserts and deletes are unbalanced (a hot ticker gets thousands of inserts on one side, very few on the other) the index becomes lopsided. A lopsided index doubles the worst-case lookup time, which is the difference between hitting the spread and missing the trade. Their alert rule is simple: at every node, the heights of the two subtrees must not differ by more than one. As soon as that property fails anywhere in any index, page the on-call DBA.

That is balanced-binary-tree. Once you see this problem in this framing, you will recognise it in every system that depends on a balanced data structure: B-trees in PostgreSQL and MySQL, AVL and red-black trees in language standard libraries, decision trees in XGBoost and LightGBM, skip lists in Redis sorted sets.

What's actually being asked

Decide whether a binary tree is height-balanced: at every node, the heights of its two subtrees differ by at most 1. Return true or false. The empty tree is balanced. A single node is balanced.

The naive approach

For every node, compute the height of its left subtree and the height of its right subtree from scratch, compare, and recurse. Each height computation is a full traversal of the subtree, so this is O(n) work per node, O(n²) total. On a skewed chain of 10⁴ nodes that is 10⁸ operations, painful but possibly within the time limit. On a slightly larger input, it dies.

The deeper problem with the naive approach is that it computes the same heights over and over. The height of the leftmost subtree is recomputed for every ancestor on the path back to the root. Wasted work.

The key insight

The whole reason the naive approach is slow is that it computes the value (height) and the predicate (balanced) as separate traversals. The trick is to return a tuple from the recursion: each call gives back both the height of the subtree AND a flag saying "this subtree is itself balanced". With both pieces of information in hand, the parent decides "I am balanced if both my children are balanced AND |hL - hR| <= 1".

One pass, one tuple per node, every node visited exactly once. O(n) instead of O(n²). This is the canonical example of why the "return a tuple" upgrade exists: a single fact (height) is not enough; a single fact (balanced) is not enough; together they are.

You can also encode "unbalanced" as a sentinel height (say, -1) instead of a real tuple. That works, but the tuple form makes the algorithm self-documenting and is the version production codebases prefer.

This is exactly the case the Trees primer mentions when it says "if you want to gather information up, return tuples".

Step-by-step approach

  1. Parse the level-order array into a tree.
  2. Define check(node):
    • If node is None, return (0, True).
    • (hL, balL) = check(node.left); if not balL, short-circuit, return (0, False).
    • (hR, balR) = check(node.right); if not balR, return (0, False).
    • If |hL - hR| > 1, return (0, False).
    • Otherwise return (1 + max(hL, hR), True).
  3. Call check(root) and print the second component.

For 10⁴-node skewed inputs, use an iterative post-order stack: push (node, False). On a False pop, push (node, True) and both children as False. On a True pop, both child heights are now known (look them up in a dict). Compute this node's height and balance flag; bail with False the moment a violation appears.

Worked example

Balanced example:

Input: [3, 9, 20, null, null, 15, 7]

       3
      / \
     9   20
        /  \
       15   7

check(9)  = (1, True)
check(15) = (1, True)
check(7)  = (1, True)
check(20): hL=1, hR=1, |0|=0 <= 1, returns (2, True)
check(3):  hL=1, hR=2, |1|=1 <= 1, returns (3, True)

Answer: True

Unbalanced example:

Input: [1, 2, 2, 3, 3, null, null, 4, 4]

         1
        / \
       2   2
      / \
     3   3
    / \
   4   4

check(4)=  (1, True), check(4) = (1, True)
check(3):  hL=1, hR=1, returns (2, True)
check(3):  no children, returns (1, True)
check(2):  hL=2, hR=1, |1| <= 1, returns (3, True)
check(2):  no children, returns (1, True)
check(1):  hL=3, hR=1, |2| > 1, returns (0, False)

Answer: False

Complexity

  • Time: O(n). Each node contributes constant work. The tuple-return upgrade is the entire reason this is not O(n²).
  • Space: O(h) for the recursion or explicit stack.

Common pitfalls

  • Naive O(n²) by separating height from balance: this is the trap. Production codebases see this regression often: junior engineers write a clean two-function solution that passes the unit tests and times out under load.
  • Re-checking balance at the root only: the spec says "for every node". A balanced root with an unbalanced internal node is invalid. The recursion must short-circuit the moment any subtree reports unbalanced.
  • Off-by-one in the height of empty: empty subtree has height 0. A single-node leaf has height 1. Pick a convention and stick with it consistently across base case and inductive step.
  • Stack overflow on a chain: 10⁴-node skewed input overflows Python's default recursion limit. Use the iterative form.

Where this shows up in the enterprise

  • PostgreSQL and MySQL B-tree indexes: when a B-tree's balance degrades (fillfactor drift, deleted-but-not-vacuumed pages), lookup latency climbs. Tools like pg_stattuple and pt-index-usage surface balance metrics.
  • MongoDB: WiredTiger storage engine indexes are B-trees; the balance metric is exposed via db.collection.stats() and watched by Atlas.
  • Redis sorted sets: skip lists are kept probabilistically balanced; degraded balance causes ZRANGEBYSCORE slowdowns and is monitored by Redis Enterprise.
  • Java's TreeMap and TreeSet: red-black trees, with rebalancing on every insert and delete. The balance invariant is exactly the spec of this problem (a relaxation of the strict balance).
  • C++ STL std::map: same red-black structure.
  • Linux kernel CFS scheduler: red-black tree of runnable tasks; balance is critical for fair scheduling.
  • XGBoost and LightGBM: gradient-boosted decision trees can grow lopsided when the loss landscape favours one feature; the training framework caps depth and rebalances after.
  • Cloudflare DNS: the underlying trie data structures are balance-checked as part of canary deployments.
  • AWS DynamoDB: range partitioning trees must stay balanced to keep p99 latency bounded.

The pattern is universal: any time a hierarchical data structure has performance guarantees, those guarantees come from a balance invariant, and somebody is paid to monitor it.

Variants you'll see elsewhere

  • count-good-nodes-in-binary-tree: count nodes for which an invariant holds along the path.
  • binary-tree-maximum-path-sum: another "return a tuple" upgrade (downward max-gain plus global max).
  • longest-path-with-different-adjacent-characters: tree DP with a tuple of (longest path through this node, longest path ending at this node).
  • AVL tree rebalancing algorithms: production version of this check, with rotations applied when the balance breaks.
  • Self-balancing BST insertion: every insertion triggers a balance check up the path.

How parikshan turns this into a learning loop

After you read this article and click through to the practice link, the parikshan flow integrates four layers:

  1. The editor runs your code against public tests and per-session dynamic private tests. A common failure mode is the naive O(n²) solution that passes the public tests and times out on a 5000-node lopsided shape in the private set. The grader's time budget is calibrated to fail the naive solution and pass the tuple-return form.
  2. Hints are graduated. The first hint asks what information you must combine at every node. The third gives you the post-order tuple recipe. Each hint costs integrity.
  3. The AI training assistant offers Explain, Review, Find Bugs, Add Comments, Optimise at small integrity costs in exam mode; free-form in practice mode.
  4. The dispute flow runs after the verdict. The common dispute on this problem is candidates who hit time limit and assume the grader is slow. Reviewers can confirm the time budget was deliberate.

Read the concept, attempt in practice mode, take hints when stuck, then redo in exam mode. When you can write the tuple-return post-order in under five minutes without touching AI, the pattern is in your fingers.

In the AI-integrated workspace

In your future job, you will not write is_balanced by hand. An agent will. Your job is to:

  1. Recognise that the balance predicate plus the height computation are two facts that must travel together; that recognition is what saves you from the O(n²) trap.
  2. Tell the agent precisely: "post-order traversal returning a (height, balanced) tuple, short-circuit on the first violation, iterative form for skewed inputs".
  3. Read the agent's code and check that it returns a tuple, not two separate function calls per node. If you see one function called height and another called is_balanced invoking it recursively, that is the O(n²) trap; send the code back.
  4. Stress-test on the edge cases the agent will silently miss: the empty tree, a single node, a right-only chain of length 2, and a tree where balance fails at a deep internal node rather than the root.

Candidates who can name "this is the tuple-return upgrade, do not separate height from balance" within five seconds are the ones who direct AI agents productively. Candidates who cannot recognise the pattern end up approving the slow version and only discover the regression when it hits production traffic. parikshan is built to train the first kind.

balanced-binary-tree