Find Leaves of Binary Tree

How it works

Instead of collecting pairs and sorting, we can directly place each node into the correct group during the DFS. When we compute a node's height, we use that height as an index into the result list. If the list doesn't have enough sublists yet, we create a new one.

Steps

1. Initialize an empty result list.
2. Perform a DFS to compute the height of each node:
   • If the node is null, return -1.
   • Recursively get the heights of left and right children.
   • Compute current height as max(leftHeight, rightHeight) + 1.
   • If the result list size equals the current height, append a new empty sublist.
   • Add the node's value to result[currentHeight].
   • Return the current height.
3. Call DFS on the root and return the result.

Complexity

• Time complexity: O(N)
• Space complexity: O(N)

Where N is the total number of nodes in the binary tree

Things to watch for

Confusing Height with Depth

Height is measured from the node down to its farthest leaf (leaves have height 0), while depth is measured from the root down to the node. This problem requires height-based grouping, not depth. Using depth would group nodes by their distance from the root, which produces incorrect results since leaves at different depths would end up in different groups.

Returning Wrong Base Case Value

The base case for null nodes should return -1, not 0. Since leaves have height 0 (computed as max(-1, -1) + 1), returning 0 for null would incorrectly make leaves have height 1. This off-by-one error shifts all groupings and produces wrong output.

Forgetting to Handle Result List Growth

When directly inserting nodes into the result list by height index, you must check if the list has enough sublists before accessing result[height]. Failing to add a new sublist when result.size() == height causes an index out of bounds error. Always ensure the result list grows dynamically as you encounter nodes with new height values.

from typing import List, Optional

class Solution:
    def findLeaves(self, root: Optional[TreeNode]) -> List[List[int]]:
        self.solution = []

        def getHeight(node):
            if node is None:
                return -1

            leftHeight = getHeight(node.left)
            rightHeight = getHeight(node.right)
            currHeight = max(leftHeight, rightHeight) + 1

            if len(self.solution) == currHeight:
                self.solution.append([])

            self.solution[currHeight].append(node.val)
            return currHeight

        getHeight(root)
        return self.solution

class Solution {
private:

    vector<vector<int>> solution;

public:

    int getHeight(TreeNode *root) {

        if (!root) {
            return -1;
        }

        int leftHeight = getHeight(root->left);
        int rightHeight = getHeight(root->right);

        int currHeight = max(leftHeight, rightHeight) + 1;

        if (this->solution.size() == currHeight) {
            this->solution.push_back({});
        }

        this->solution[currHeight].push_back(root->val);

        return currHeight;
    }

    vector<vector<int>> findLeaves(TreeNode* root) {
        this->solution.clear();

        getHeight(root);

        return this->solution;
    }
};