Find Leaves of Binary Tree — LC 366 Height-Based Grouping

Sanjeev SharmaSanjeev Sharma
5 min read

Advertisement

Problem Statement

LeetCode 366 — Find Leaves of Binary Tree | Difficulty: Medium

Given the root of a binary tree, collect the tree's nodes as if you were doing this: collect all leaf nodes, remove them, and repeat until the tree is empty. Return the collected nodes at each step.

Constraints:

  • The number of nodes is in the range [1, 100]
  • -100 <= Node.val <= 100
Input:  root = [1,2,3,4,5]
Output: [[4,5,3],[2],[1]]
Input:  root = [1]
Output: [[1]]

Why This Problem Matters

Find Leaves of Binary Tree is a LinkedIn and Amazon interview problem that looks straightforward but rewards a non-obvious insight: instead of actually removing leaves repeatedly (which would require multiple passes and tree modification), you can determine every node's "leaf round" in a single O(n) DFS by computing each node's height — its distance from the nearest leaf below it.

This height-based grouping is a reusable pattern: any problem that asks you to process nodes "in order of some property computed bottom-up" can use this approach. It also tests whether candidates know when to abandon the obvious simulation in favor of a smarter key-function computation.

The Core Insight

A node becomes a leaf in round h where h is its height in the tree:

  • Leaf nodes have height 0 (collected in round 1)
  • Their parents have height 1 (collected in round 2 after leaves are removed)
  • And so on upward

Height is computed bottom-up (post-order):

  • Null nodes have height -1
  • All other nodes: height = 1 + max(height_left, height_right)

Group nodes by their computed height. Nodes at height 0 form the first result group, height 1 the second, and so on. No actual tree modification needed.

Visual Dry Run

Tree: [1, 2, 3, 4, 5]

Nodeleft heightright heightnode heightgroup
4 (leaf)-1-10result[0]
5 (leaf)-1-10result[0]
3 (leaf)-1-10result[0]
2001result[1]
1102result[2]

Result: [[4, 5, 3], [2], [1]]

Note: order within each group depends on DFS traversal order (left-right post-order by default).

Solution (Optimal)

class Solution:
    def findLeaves(self, root):
        result = []
 
        def dfs(node):
            if not node:
                return -1  # null node height
 
            # Post-order: compute children heights first
            left_h = dfs(node.left)
            right_h = dfs(node.right)
 
            # Height of this node
            h = 1 + max(left_h, right_h)
 
            # Extend result list if needed
            if len(result) == h:
                result.append([])
 
            result[h].append(node.val)
            return h
 
        dfs(root)
        return result
var findLeaves = function(root) {
    const result = [];
 
    function dfs(node) {
        if (!node) return -1;
 
        const leftH = dfs(node.left);
        const rightH = dfs(node.right);
 
        const h = 1 + Math.max(leftH, rightH);
 
        if (result.length === h) result.push([]);
        result[h].push(node.val);
 
        return h;
    }
 
    dfs(root);
    return result;
};

Time: O(n) — each node visited exactly once Space: O(n) — result array stores all n nodes; call stack is O(h)

Common Mistakes

  • Simulating the process by actually removing leaves and re-running DFS — this is O(n^2) in the worst case and modifies the input tree
  • Confusing node height (distance from leaf) with node depth (distance from root) — this problem uses height, not depth
  • Initializing null height to 0 instead of -1 — causes leaf nodes to have height 1 instead of 0, shifting all groups by one
  • Using result[h] = [] without checking if the index exists — use append/push when h equals the current length

Interview Tips

  • State the key insight early: "I compute each node's height and group by height — no tree modification needed"
  • Distinguish between height (distance from leaf) and depth (distance from root) clearly
  • The if len(result) == h: result.append([]) idiom ensures the result array grows exactly as needed
  • A single DFS pass collects everything — mention this O(n) vs O(n^2) comparison

Follow-up Questions

  • What if you need to actually remove leaves from the tree? Set node.left = None and node.right = None after computing height — the tree is structurally modified as desired.
  • What if the tree is n-ary? Change max(height_left, height_right) to max(heights of all children). The rest of the code is identical.
  • How does this differ from BFS level-order traversal? Level-order groups by depth (distance from root). This problem groups by height (distance from nearest leaf), which requires post-order DFS.
  • What if you need the groups in reverse order (root first)? Return result[::-1] at the end, or compute depth instead of height.

Key Takeaways

  • Nodes belong to the same "leaf round" if and only if they have the same height (distance from nearest leaf)
  • Height computed post-order: h = 1 + max(left_h, right_h); null nodes have height -1
  • Grouping by height in a single DFS pass is O(n) — much better than repeated leaf removal at O(n^2)
  • The result array grows dynamically: when len(result) == h, append a new empty list
  • No tree modification needed — purely observational classification by height
  • This height-based grouping pattern applies to any problem asking "in what order would nodes be processed bottom-up?"
  • Time O(n), space O(n) — optimal since all nodes must appear in the result

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro

Related reading