Find Leaves of Binary Tree

LeetCode 366: Find Leaves of Binary Tree Solution & Explanation

Problem Overview

TL;DR: The optimal solution utilizes a postorder traversal to calculate the height of each node, grouping nodes with the same height into the same output list.

The Find Leaves of Binary Tree problem requires us to dismantle a binary tree layer by layer. Starting from the outside, we must collect all leaf nodes, remove them, and repeat the process until the tree is empty. The output must be a structured list of lists, where the first list contains the initial leaves, the second list contains the new leaves formed after the first removal, and so on.

This is a popular interview question that tests your ability to manipulate tree data structures and understand traversal orders. While the problem description suggests physically modifying the tree structure, the LeetCode 366 solution is best implemented by analyzing node properties without expensive structural modifications.

Brute Force Approach for Find Leaves of Binary Tree

The most intuitive way to solve this problem is to simulate the process exactly as described. We can write a function that traverses the tree, identifies leaf nodes, adds them to a current list, and disconnects them from their parents. We repeat this traversal until the root is null.

Pseudo-code:

while root is not null:
    current_leaves = []
    root = remove_leaves(root, current_leaves)
    result.add(current_leaves)
return result

function remove_leaves(node, list):
    if node is null: return null
    if node is leaf:
        list.add(node.val)
        return null // Remove the node
    node.left = remove_leaves(node.left, list)
    node.right = remove_leaves(node.right, list)
    return node

Time Complexity: $O(N^2)$. In the worst case (a skewed tree like a linked list), we traverse the entire remaining height of the tree for every single node removal. Why it fails: This approach is inefficient for large trees or strict time constraints. It visits nodes near the top of the tree repeatedly. We need a solution that visits every node exactly once ($O(N)$).

Algorithm Strategy: Tree DFS - Recursive Postorder Traversal

To implement the optimal LeetCode 366 solution, we follow these steps using the Postorder Traversal pattern:

1. Initialize Result Structure: Create a list of lists (e.g., vector<vector<int>> or List<List<Integer>>) to store the values.
2. Helper Function (DFS): Define a recursive function that takes a node as input and returns an integer representing the node's height.
3. Base Case: If the current node is null, return -1. This allows leaf nodes to calculate their height as $1 + \max(-1, -1) = 0$.
4. Recursive Step (Postorder):
   • Recursively call the helper on the left child to get leftHeight.
   • Recursively call the helper on the right child to get rightHeight.
5. Logic:
   • Calculate currentHeight = 1 + max(leftHeight, rightHeight).
   • Ensure the result list has a size large enough to accommodate the index currentHeight. If the size is equal to currentHeight, add a new empty list.
   • Add the current node's value to result[currentHeight].
6. Return: Return currentHeight to the parent caller.

This strategy ensures every node is visited exactly once, and values are sorted into rounds purely based on their topological distance from the leaves.

Execution Flow

Let's trace the execution on a simple tree: [1, 2, 3] (1 is root, 2 is left child, 3 is right child).

1. Start at Root (1): Call getHeight(1).
2. Traverse Left: Call getHeight(2).
   • Call getHeight(2.left) -> returns -1 (Base case).
   • Call getHeight(2.right) -> returns -1 (Base case).
   • Calculate height of 2: 1 + max(-1, -1) = 0.
   • result size is 0. Add new list. result is [[]].
   • Add 2 to result[0]. result is [[2]].
   • Return 0.
3. Traverse Right: Call getHeight(3).
   • Call getHeight(3.left) -> returns -1.
   • Call getHeight(3.right) -> returns -1.
   • Calculate height of 3: 1 + max(-1, -1) = 0.
   • result size is 1. 0 < 1, so no resize needed.
   • Add 3 to result[0]. result is [[2, 3]].
   • Return 0.
4. Process Root (1):
   • Received leftHeight = 0 from node 2.
   • Received rightHeight = 0 from node 3.
   • Calculate height of 1: 1 + max(0, 0) = 1.
   • result size is 1. 1 == 1, so add new list. result is [[2, 3], []].
   • Add 1 to result[1]. result is [[2, 3], [1]].
   • Return 1.
5. Final Output: [[2, 3], [1]].

Proof of Correctness

The correctness relies on the definition of the problem: a node is removed only after all its descendants are removed.

• Base Case: A leaf has no descendants. Its "removal dependency chain" length is 0. Our algorithm assigns it height 0.
• Inductive Step: For any node $N$, it can only be removed after its last remaining child is removed. If the left child takes $L$ rounds to be removed and the right child takes $R$ rounds, node $N$ must wait for $\max(L, R)$ rounds to complete. Therefore, node $N$ is removed in round $\max(L, R) + 1$. This strictly matches the height calculation logic: $Height(N) = 1 + \max(Height(Left), Height(Right))$. Thus, grouping by height is equivalent to grouping by removal round.

Pattern Reuse Notes

The Tree DFS - Recursive Postorder Traversal pattern is fundamental for problems involving node heights or bottom-up aggregation.

• LeetCode 104: Maximum Depth of Binary Tree: Uses the exact same logic (1 + max(left, right)) to find the total height of the root.
• LeetCode 110: Balanced Binary Tree: Uses postorder traversal to calculate heights and check the balance condition simultaneously.
• LeetCode 124: Binary Tree Maximum Path Sum: Bubbles up the maximum path gain from children to compute the optimal path through the current node.
• LeetCode 145: Binary Tree Postorder Traversal: The fundamental traversal mechanism used in all these solutions.

Reference Implementation

C++ Solution for LeetCode 366

1class Solution {
2public:
3    vector<vector<int>> findLeaves(TreeNode* root) {
4        vector<vector<int>> result;
5        getHeight(root, result);
6        return result;
7    }
8
9private:
10    int getHeight(TreeNode* node, vector<vector<int>>& result) {
11        // Base case: null nodes return -1 height
12        if (!node) {
13            return -1;
14        }
15
16        // Postorder traversal: process children first
17        int leftHeight = getHeight(node->left, result);
18        int rightHeight = getHeight(node->right, result);
19
20        // Current height is 1 + max height of children
21        int currHeight = 1 + max(leftHeight, rightHeight);
22
23        // Ensure the result vector is large enough for this height level
24        if (result.size() == currHeight) {
25            result.push_back({});
26        }
27
28        // Add current node value to the corresponding height list
29        result[currHeight].push_back(node->val);
30
31        return currHeight;
32    }
33};

Java Solution for LeetCode 366

1class Solution {
2    public List<List<Integer>> findLeaves(TreeNode root) {
3        List<List<Integer>> result = new ArrayList<>();
4        getHeight(root, result);
5        return result;
6    }
7
8    private int getHeight(TreeNode node, List<List<Integer>> result) {
9        // Base case: null nodes return -1 height
10        if (node == null) {
11            return -1;
12        }
13
14        // Postorder traversal: process children first
15        int leftHeight = getHeight(node.left, result);
16        int rightHeight = getHeight(node.right, result);
17
18        // Current height is 1 + max height of children
19        int currHeight = 1 + Math.max(leftHeight, rightHeight);
20
21        // Ensure the result list is large enough for this height level
22        if (result.size() == currHeight) {
23            result.add(new ArrayList<>());
24        }
25
26        // Add current node value to the corresponding height list
27        result.get(currHeight).add(node.val);
28
29        return currHeight;
30    }
31}

Python Solution for LeetCode 366

1class Solution:
2    def findLeaves(self, root: Optional[TreeNode]) -> List[List[int]]:
3        result = []
4        
5        def get_height(node):
6            # Base case: null nodes return -1 height
7            if not node:
8                return -1
9            
10            # Postorder traversal: process children first
11            left_height = get_height(node.left)
12            right_height = get_height(node.right)
13            
14            # Current height is 1 + max height of children
15            curr_height = 1 + max(left_height, right_height)
16            
17            # Ensure the result list is large enough for this height level
18            if len(result) == curr_height:
19                result.append([])
20            
21            # Add current node value to the corresponding height list
22            result[curr_height].append(node.val)
23            
24            return curr_height
25            
26        get_height(root)
27        return result