Problem Description

You are given the root of a binary tree. Your task is to collect the nodes of the tree in a specific way that simulates repeatedly removing leaf nodes.

The process works as follows:

1. First, identify and collect all current leaf nodes (nodes with no children) into a group
2. Remove these leaf nodes from the tree
3. After removal, some nodes that previously had children may now become new leaf nodes
4. Repeat steps 1-2, collecting each new set of leaf nodes into separate groups
5. Continue this process until the entire tree is empty

The result should be a list of lists, where each inner list contains the values of nodes that were removed together in the same iteration.

For example, if you have a tree like:

      1
     / \
    2   3
   / \
  4   5

The collection process would be:

• First iteration: Collect leaves [4, 5, 3] and remove them
• Second iteration: Now 2 is a leaf, collect [2] and remove it
• Third iteration: Now 1 is a leaf, collect [1] and remove it

The final output would be: [[4, 5, 3], [2], [1]]

The solution uses a depth-first search (DFS) approach that calculates the "height" of each node from the bottom. Nodes at the same height from the bottom will be removed in the same iteration. Leaf nodes have height 0, their parents have height 1, and so on. The algorithm groups nodes by their height, which naturally corresponds to the order in which they would be removed as leaves.

Flowchart Walkthrough

First, let's pin down the algorithm using the Flowchart. Here's a step-by-step walkthrough:

Is it a graph?

• Yes: A binary tree is a special type of graph where each node has at most two children and there are no cycles.

Is it a tree?

• Yes: The problem explicitly states we're working with a binary tree structure.

DFS

• Following the "yes" path from "Is it a tree?", we arrive at DFS (Depth-First Search).

Conclusion: The flowchart suggests using a DFS approach for this problem.

This aligns perfectly with the solution approach. The problem requires us to process nodes based on their distance from the leaf nodes (height from bottom). DFS is ideal here because:

1. We need to traverse to the leaves first to determine their height (0 for leaves)
2. We then work our way back up, calculating each node's height based on its children's heights
3. The post-order traversal nature of DFS allows us to process children before their parents, which is exactly what we need to calculate heights from bottom to top

The DFS solution cleverly groups nodes by their height from the bottom, where:

• Height 0 = current leaf nodes (removed first)
• Height 1 = nodes that become leaves after first removal
• Height 2 = nodes that become leaves after second removal
• And so on...

This height-based grouping naturally gives us the required output format without actually simulating the removal process.

Intuition

The key insight is recognizing that we don't actually need to simulate the removal process. Instead, we can think about when each node would be removed based on its position in the tree.

Consider what makes a node removable at each step - it must be a leaf node. Initially, only the actual leaf nodes can be removed. After removing them, their parent nodes might become new leaves if they have no other children. This pattern continues until we reach the root.

This observation leads to an important realization: the order in which nodes are removed depends on their distance from the nearest leaf node. We can define this distance as the "height from bottom":

• Actual leaf nodes have height 0
• Parents of leaf nodes have height 1 (they become leaves after one removal)
• Grandparents of leaf nodes have height 2 (they become leaves after two removals)
• And so on...

Instead of repeatedly finding and removing leaves (which would be inefficient), we can calculate each node's height from the bottom in a single DFS traversal. The height of any node is simply 1 + max(height of left child, height of right child). For null nodes, we return height -1, so leaf nodes correctly get height 0.

During the DFS traversal, as we calculate each node's height, we can directly place it into the appropriate group in our result. All nodes with the same height will be removed together in the same iteration, so they belong in the same sublist.

This transforms the problem from "simulate leaf removal" to "group nodes by their height from bottom", which can be solved efficiently with a single post-order DFS traversal where we process children before parents to calculate heights bottom-up.

Solution Approach

The solution implements a post-order DFS traversal to calculate the height of each node from the bottom and group nodes by their heights.

Data Structure:

• ans: A list of lists where ans[h] contains all nodes at height h from the bottom

Algorithm Steps:

1. Define the DFS function that returns the height of each node:

   def dfs(root: Optional[TreeNode]) -> int:

2. Base case: If the node is None, return 0. This ensures leaf nodes will have height 0 after calculation:

   if root is None:
       return 0

3. Recursive calls: Calculate heights of left and right subtrees:

   l, r = dfs(root.left), dfs(root.right)
   h = max(l, r)

   The current node's height index is the maximum of its children's heights.

4. Dynamic list expansion: If we encounter a new height level, add a new list to accommodate it:

   if len(ans) == h:
       ans.append([])

5. Group nodes by height: Add the current node's value to the appropriate height group:

   ans[h].append(root.val)

6. Return height for parent: Return h + 1 so the parent node knows this subtree's height:

   return h + 1

Why this works:

• Leaf nodes have no children, so dfs(left) and dfs(right) both return 0, making h = max(0, 0) = 0. They're placed in ans[0].
• Parents of leaves get h = max(child_heights) = 0, so they're placed in ans[0], but return 1 to their parents.
• This pattern continues, with each node being placed at index equal to its maximum child's height, creating the correct grouping.

Time Complexity: O(n) where n is the number of nodes, as we visit each node exactly once.

Space Complexity: O(n) for the recursion stack in the worst case (skewed tree) and for storing the result.

Example Walkthrough

Let's trace through a small example to illustrate the solution approach:

    1
   / \
  2   3
 /
4

Step 1: Start DFS from root (node 1)

• Call dfs(1), which needs to process its children first (post-order)

Step 2: Process left subtree of node 1

• Call dfs(2), which needs to process its children first
• Call dfs(4), which has no children
  • dfs(4.left) returns 0 (null child)
  • dfs(4.right) returns 0 (null child)
  • h = max(0, 0) = 0
  • Add 4 to ans[0]: ans = [[4]]
  • Return 0 + 1 = 1 to parent

Step 3: Complete processing node 2

• Left child (4) returned 1, right child (null) returns 0
• h = max(1, 0) = 1
• Add 2 to ans[1]: ans = [[4], [2]]
• Return 1 + 1 = 2 to parent

Step 4: Process right subtree of node 1

• Call dfs(3), which has no children
  • Both children return 0 (null)
  • h = max(0, 0) = 0
  • Add 3 to ans[0]: ans = [[4, 3], [2]]
  • Return 0 + 1 = 1 to parent

Step 5: Complete processing root (node 1)

• Left child (2) returned 2, right child (3) returned 1
• h = max(2, 1) = 2
• Add 1 to ans[2]: ans = [[4, 3], [2], [1]]
• Return 2 + 1 = 3

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

This matches the expected leaf removal order:

• First removal: leaves 4 and 3
• Second removal: node 2 becomes a leaf after 4 is removed
• Third removal: node 1 becomes a leaf after 2 and 3 are removed

Solution Implementation

Time and Space Complexity

Time Complexity: O(n) where n is the number of nodes in the binary tree.

The algorithm performs a depth-first search (DFS) traversal of the tree, visiting each node exactly once. At each node, the operations performed are:

• Recursive calls to left and right children
• Finding the maximum of two values: max(l, r) - O(1)
• Checking list length and potentially appending a new list: O(1) amortized
• Appending the node value to a list: O(1) amortized
• Returning a value: O(1)

Since each node is visited once and each visit performs O(1) operations, the total time complexity is O(n).

Space Complexity: O(n) where n is the number of nodes in the binary tree.

The space complexity consists of:

• Recursive call stack: In the worst case (skewed tree), the recursion depth can be O(n). In a balanced tree, it would be O(log n). The worst-case is O(n).
• Output storage (ans list): The result stores all n node values across multiple sublists, requiring O(n) space.
• Local variables: Each recursive call uses constant space for variables l, r, and h, but this is already accounted for in the call stack analysis.

The dominant factors are the recursive call stack and the output storage, both potentially requiring O(n) space, resulting in an overall space complexity of O(n).

Common Pitfalls

1. Incorrect Height Calculation for None Nodes

Pitfall: Returning 0 for None nodes instead of -1, which causes leaf nodes to have height 1 instead of 0.

# WRONG - This makes leaf nodes have height 1
def dfs(node):
    if node is None:
        return 0  # Incorrect!
    left_height = dfs(node.left)
    right_height = dfs(node.right)
    current_height = max(left_height, right_height) + 1
    # Leaf nodes will have height = max(0, 0) + 1 = 1

Solution: Return -1 for None nodes so that leaf nodes correctly get height 0:

# CORRECT
def dfs(node):
    if node is None:
        return -1  # Correct!
    left_height = dfs(node.left)
    right_height = dfs(node.right)
    current_height = max(left_height, right_height) + 1
    # Leaf nodes will have height = max(-1, -1) + 1 = 0

2. Forgetting to Dynamically Expand the Result List

Pitfall: Assuming the result list already has enough sublists, leading to IndexError:

# WRONG - Will crash with IndexError
def dfs(node):
    if node is None:
        return -1
    left_height = dfs(node.left)
    right_height = dfs(node.right)
    current_height = max(left_height, right_height) + 1
    result[current_height].append(node.val)  # IndexError if result[current_height] doesn't exist!
    return current_height

Solution: Check and expand the result list before accessing:

# CORRECT
def dfs(node):
    if node is None:
        return -1
    left_height = dfs(node.left)
    right_height = dfs(node.right)
    current_height = max(left_height, right_height) + 1
  
    # Ensure the list has enough sublists
    if len(result) == current_height:
        result.append([])
  
    result[current_height].append(node.val)
    return current_height

3. Using Pre-order Instead of Post-order Traversal

Pitfall: Processing the current node before its children, which means you don't know the node's height yet:

# WRONG - Pre-order traversal doesn't work
def dfs(node):
    if node is None:
        return -1
  
    # Can't determine height here without visiting children first!
    result[???].append(node.val)  # What height to use?
  
    left_height = dfs(node.left)
    right_height = dfs(node.right)
    return max(left_height, right_height) + 1

Solution: Use post-order traversal - process children first, then the current node:

# CORRECT - Post-order traversal
def dfs(node):
    if node is None:
        return -1
  
    # Process children first
    left_height = dfs(node.left)
    right_height = dfs(node.right)
  
    # Now we know the height and can process current node
    current_height = max(left_height, right_height) + 1
    if len(result) == current_height:
        result.append([])
    result[current_height].append(node.val)
  
    return current_height