Code Implementation

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def findLeaves(self, root):
        res = []
        def dfs(node):
            if not node:
                return -1
            left = dfs(node.left)
            right = dfs(node.right)
            level = max(left, right) + 1
            if level == len(res):
                res.append([])
            res[level].append(node.val)
            return level
        dfs(root)
        return res
      

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    vector<vector<int>> findLeaves(TreeNode* root) {
        vector<vector<int>> res;
        function<int(TreeNode*)> dfs = [&](TreeNode* node) {
            if (!node) return -1;
            int left = dfs(node->left);
            int right = dfs(node->right);
            int level = max(left, right) + 1;
            if (level == res.size())
                res.push_back({});
            res[level].push_back(node->val);
            return level;
        };
        dfs(root);
        return res;
    }
};
      

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    public List<List<Integer>> findLeaves(TreeNode root) {
        List<List<Integer>> res = new ArrayList<>();
        dfs(root, res);
        return res;
    }
    private int dfs(TreeNode node, List<List<Integer>> res) {
        if (node == null) return -1;
        int left = dfs(node.left, res);
        int right = dfs(node.right, res);
        int level = Math.max(left, right) + 1;
        if (level == res.size())
            res.add(new ArrayList<>());
        res.get(level).add(node.val);
        return level;
    }
}
      

/**
 * Definition for a binary tree node.
 * function TreeNode(val, left, right) {
 *     this.val = (val===undefined ? 0 : val)
 *     this.left = (left===undefined ? null : left)
 *     this.right = (right===undefined ? null : right)
 * }
 */
/**
 * @param {TreeNode} root
 * @return {number[][]}
 */
var findLeaves = function(root) {
    const res = [];
    function dfs(node) {
        if (!node) return -1;
        let left = dfs(node.left);
        let right = dfs(node.right);
        let level = Math.max(left, right) + 1;
        if (level === res.length) res.push([]);
        res[level].push(node.val);
        return level;
    }
    dfs(root);
    return res;
};
      

Problem Description

Given the root of a binary tree, repeatedly remove all of its leaves, and collect the leaves removed at each step until the tree is empty. Return a list of lists, where each inner list contains the values of the leaves removed at that step.

• The binary tree is represented using TreeNode objects with val, left, and right attributes.
• Each time, remove all current leaves (nodes with no children) and collect their values together.
• Repeat the process until the tree is empty.
• The output should be a list of lists, with each inner list containing the leaf values removed at that round.
• Each node will appear in exactly one inner list, and you must not reuse nodes.

Thought Process

At first glance, the problem seems to require simulating the process of removing leaves from a binary tree, layer by layer. A brute-force approach would be to physically remove all leaves in each iteration, traverse the tree to find the leaves, and repeat until the tree is empty. However, this would be inefficient, as it would require multiple traversals and modifications of the tree.

Instead, we can notice that each node is "removed" as a leaf at a certain step, which depends on its distance from the farthest leaf. Specifically, a node becomes a leaf after all of its children have been removed. This leads to the insight that we can assign a "level" to each node: leaves are level 0, their parents are level 1, and so on.

By performing a post-order traversal, we can compute the level for each node in a single pass, grouping nodes by their level to simulate the rounds of leaf removal.

Solution Approach

• Step 1: Post-Order Traversal
  • Traverse the tree using post-order (left, right, node).
  • For each node, recursively compute the level of its left and right children.
• Step 2: Compute Node Level
  • The level of a node is max(left child level, right child level) + 1.
  • Leaves (nodes with no children) will have level 0.
• Step 3: Group Nodes by Level
  • Use a list of lists (or array of arrays) to collect node values by their computed level.
  • If there is no list for the current level, create one.
• Step 4: Return the Result
  • After traversing the entire tree, return the list of lists, where each inner list contains the values of nodes removed at that round.

This approach does not modify the tree and collects all results in a single traversal, making it efficient and elegant.

Example Walkthrough

Consider the following binary tree:

        1
       / \
      2   3
     / \
    4   5
  

Let's trace the process:

1. First round: Leaves are 4, 5, and 3. Remove them and collect [4, 5, 3].
2. Second round: Now, 2 becomes a leaf (since its children are gone). Remove and collect [2].
3. Third round: Now, 1 becomes a leaf. Remove and collect [1].

The final output is [[4, 5, 3], [2], [1]].

Using the algorithm, each node is assigned a level: 4, 5, and 3 get level 0; 2 gets level 1; 1 gets level 2. Grouping by level gives the correct answer.

Time and Space Complexity

• Brute-force approach:
  • Each round requires traversing the tree to find and remove leaves, with up to O(n) rounds and O(n) traversal per round, leading to O(n²) time complexity.
  • Space complexity is O(n) for storing results.
• Optimized approach (current solution):
  • Each node is visited exactly once in a post-order traversal: O(n) time.
  • Space complexity is O(n) for the result and call stack (in the worst case, e.g., a skewed tree).

Summary

By rethinking the problem as grouping nodes by the step at which they become leaves, we avoid repeatedly modifying the tree and instead use a single post-order traversal to assign levels to nodes. This results in an efficient O(n) time and space solution that is easy to implement and understand. The key insight is that each node's "removal round" is determined by its maximum distance from a leaf, enabling us to build the answer in one pass.