Taro Logo

Find Leaves of Binary Tree

Medium
LinkedIn logo
LinkedIn
2 views
Topics:
TreesRecursion

Given the root of a binary tree, collect a tree's nodes as if you were doing this:

  • Collect all the leaf nodes.
  • Remove all the leaf nodes.
  • Repeat until the tree is empty.

Return a list of lists of the collected nodes.

Example 1:

Input: root = [1,2,3,4,5]
Output: [[4,5,3],[2],[1]]
Explanation:
[[4,5,3],[2],[1]]

Example 2:

Input: root = [1]
Output: [[1]]

Constraints:

  • The number of nodes in the tree is in the range [1, 100].
  • -100 <= Node.val <= 100

Solution


Clarifying Questions

When you get asked this question in a real-life environment, it will often be ambiguous (especially at FAANG). Make sure to ask these questions in that case:

  1. Can the binary tree be empty? If so, what should I return?
  2. What is the range of values for the nodes in the binary tree? Are negative values possible?
  3. If after removing leaves, the root itself becomes a leaf, should it be included in the list of removed leaves for that iteration?
  4. Is the provided tree guaranteed to be a valid binary tree (no cycles, each node has at most two children)?
  5. If the tree consists of only one node (the root), what should the output be?

Brute Force Solution

Approach

Think of peeling off the leaves from a tree, layer by layer. The brute force approach focuses on repeatedly identifying and removing all the leaves from the tree until nothing is left. We collect the leaves at each stage, and that gives us our answer.

Here's how the algorithm would work step-by-step:

  1. First, find all the nodes in the tree that are leaves. A leaf is a node that has no children.
  2. Collect these leaves into a group.
  3. Now, remove all these leaf nodes from the tree. This might involve updating the 'parent' nodes to no longer point to these removed leaf nodes.
  4. Repeat the process: find the new leaf nodes in the updated tree (the nodes that were previously parents of the old leaves).
  5. Collect these new leaves into another group.
  6. Again, remove these new leaf nodes from the tree.
  7. Continue repeating this process of finding leaves, collecting them, and removing them, until the entire tree is empty.
  8. The groups of leaves you collected at each step represent the layers of leaves you peeled off the tree.

Code Implementation

def find_leaves_of_binary_tree(root):
    leaf_layers = []

    while root:
        current_leaves = []
        new_root = remove_leaves(root, current_leaves)

        leaf_layers.append(current_leaves)
        root = new_root

    return leaf_layers

def remove_leaves(root, current_leaves):
    if not root:
        return None

    if not root.left and not root.right:
        current_leaves.append(root.val)
        return None

    # Recursively remove leaves from the left subtree
    root.left = remove_leaves(root.left, current_leaves)

    # Recursively remove leaves from the right subtree
    root.right = remove_leaves(root.right, current_leaves)

    return root

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

Big(O) Analysis

Time Complexity
O(n²)The algorithm iteratively removes leaves from the tree. In the worst-case scenario, each node can be a leaf in some iteration. Finding leaves in each iteration requires traversing the tree, which takes O(n) time in the worst case (e.g., a skewed tree where finding the next leaf requires visiting almost all nodes). Since we repeat this process of finding and removing leaves until the tree is empty, and in the worst case, we perform n iterations (where n is the number of nodes), the overall time complexity becomes O(n * n), or O(n²).
Space Complexity
O(N)The algorithm stores leaves in a list in each iteration, and ultimately stores all these lists of leaves. In the worst-case scenario (e.g., a skewed tree), each node could be a leaf at some stage, leading to the storage of almost all nodes in these leaf lists. Thus, the space required to store these lists of leaves can grow linearly with the number of nodes, N, in the tree. Therefore, the auxiliary space complexity is O(N).

Optimal Solution

Approach

The goal is to repeatedly find and remove the leaves (nodes with no children) of a binary tree, collecting them into lists. We can figure out how 'deep' each node is, working from the bottom up, which tells us when it will be a leaf.

Here's how the algorithm would work step-by-step:

  1. Think of the height of each node in the tree. Leaves have a height of 0. The parent of leaves has a height of 1, and so on.
  2. Calculate the height of each node in the tree. We do this from the bottom up, so start at the leaves and move towards the root. A node's height is one more than the maximum height of its children.
  3. Use the height of each node as the index to store the node in a list of lists. All nodes with the same height are grouped in the same list.
  4. Because we found the heights correctly, each list represents the leaves at that 'level'. For example, the list at index 0 contains all the original leaves, the list at index 1 contains the nodes that became leaves after the first removal, and so on.
  5. The list of lists is the answer, where each inner list represents the leaves at a specific stage of removal.

Code Implementation

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def find_leaves(root):
    leaves_by_height = []

    def get_height(node):
        if not node:
            return -1

        left_height = get_height(node.left)
        right_height = get_height(node.right)

        # Node height is one greater than max child height.
        node_height = max(left_height, right_height) + 1

        # Add new height list if needed
        if len(leaves_by_height) <= node_height:
            leaves_by_height.append([])

        # Store nodes at their corresponding height.
        leaves_by_height[node_height].append(node.val)
        return node_height

    get_height(root)

    return leaves_by_height

Big(O) Analysis

Time Complexity
O(n)The algorithm visits each node in the binary tree exactly once to calculate its height. Calculating the height involves constant-time operations (comparing heights of children and adding 1). Since the number of nodes in the tree is 'n', the total time complexity is directly proportional to 'n'. Therefore, the overall time complexity is O(n).
Space Complexity
O(N)The algorithm uses a list of lists to store the leaves at each level. In the worst-case scenario, where the tree is a skewed tree (e.g., a linked list), the height of the tree is N, where N is the number of nodes. Therefore, the outer list will have a size of N, and the inner lists combined will contain all N nodes. Hence, the auxiliary space used is proportional to N, resulting in O(N) space complexity.

Edge Cases

CaseHow to Handle
Null or empty treeReturn an empty list immediately as there are no nodes to process.
Single node tree (root is also a leaf)Return a list containing a single inner list with the root's value; the next iteration will result in an empty tree.
Tree with all nodes having the same valueThe algorithm should correctly identify and remove leaves regardless of value duplication.
Completely skewed tree (left or right leaning)The solution should handle the potentially deep recursion stack without stack overflow, potentially requiring tail-call optimization or iterative solution.
Tree with a large number of nodesEnsure the solution is efficient in terms of time and space complexity to avoid exceeding time limits or memory constraints, potentially using iterative removal.
Tree with negative node valuesThe algorithm should function correctly with negative node values as node values are treated as data, not indices.
A tree where some subtrees become empty before othersThe height calculation or similar logic in recursive solutions should properly account for subtrees already processed and removed.
Extreme node values that could lead to integer overflow when summing or comparing in helper functionsUse appropriate data types (e.g., long) to avoid integer overflow if operations on node values are performed.