deepSWE.io

Binary Tree Pruning

Trees

Problem Statement

Given the `root` of a binary tree, return the same tree where every subtree (of the given tree) not containing a `1` has been removed.

Example

Output: [1,null,1,null,1]

Algorithm Explanation

We can solve this problem using a post-order traversal (DFS). The idea is to decide whether to prune a subtree after visiting its children.

Algorithm Steps

Post-order Traversal: We'll use a recursive function that returns whether the subtree at the current node contains a `1`.
Recursive Calls: First, we make recursive calls on the left and right children. If a call returns `false`, it means that subtree does not contain a `1`, so we can prune it by setting the corresponding child pointer to `null`.
Return Value: The function should return `true` if the current node's value is `1` or if either of its subtrees (before pruning) contained a `1`. Otherwise, it returns `false`.

Start

Starting pruning process.

Binary Tree Pruning Solution


class Solution:
    def pruneTree(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
        if not root:
            return None
        
        root.left = self.pruneTree(root.left)
        root.right = self.pruneTree(root.right)
        
        if not root.left and not root.right and root.val == 0:
            return None
            
        return root