deepSWE.io

Find Leaves of a Binary Tree

TreesDFS

Problem Statement

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.

Example

Output: [[4, 5, 3], [2], [1]]

Algorithm Explanation

We can solve this problem using a depth-first search (DFS) approach. The key idea is to use the height of a node to determine which list of leaves it belongs to. The height of a node is the number of edges on the longest path from the node to a leaf.

Algorithm Steps

DFS with Height: We'll create a recursive DFS function that returns the height of the current node.
Base Case: If the node is null, its height is -1.
Calculate Height: The height of a node is `1 + max(height(left_child), height(right_child))`.
Store Leaves: The height of a node determines which list of leaves it belongs to. We can use a list of lists (or a map) to store the leaves. If the height of a node is `h`, we add it to the list at index `h`.

Start

Starting DFS from the root.

Result: []

Find Leaves of a Binary Tree Solution


class Solution:
    def findLeaves(self, root: Optional[TreeNode]) -> List[List[int]]:
        result = []
        
        def get_height(node):
            if not node:
                return -1
            
            left_height = get_height(node.left)
            right_height = get_height(node.right)
            
            height = 1 + max(left_height, right_height)
            
            if height == len(result):
                result.append([])
            
            result[height].append(node.val)
            return height

        get_height(root)
        return result