Binary Tree Pruning

medium trees dfs recursion

Problem

Given the root of a binary tree containing 0s and 1s, return the same tree with every subtree not containing any 1 removed.

Inputroot = [1,null,0,0,1]

Output[1,null,0,null,1]

The 0-only left subtree of the right child is pruned.

root

def pruneTree(root):
    if root is None:
        return None
    root.left = pruneTree(root.left)
    root.right = pruneTree(root.right)
    if root.val == 0 and root.left is None and root.right is None:
        return None
    return root

var pruneTree = function(root) {
    if (!root) return null;
    root.left = pruneTree(root.left);
    root.right = pruneTree(root.right);
    if (root.val === 0 && !root.left && !root.right) return null;
    return root;
};

class Solution {
    public TreeNode pruneTree(TreeNode root) {
        if (root == null) return null;
        root.left = pruneTree(root.left);
        root.right = pruneTree(root.right);
        if (root.val == 0 && root.left == null && root.right == null) return null;
        return root;
    }
}

class Solution {
public:
    TreeNode* pruneTree(TreeNode* root) {
        if (!root) return nullptr;
        root->left = pruneTree(root->left);
        root->right = pruneTree(root->right);
        if (root->val == 0 && !root->left && !root->right) return nullptr;
        return root;
    }
};

Explanation

We need to delete any subtree that holds no 1 at all. The elegant way is a post-order recursion: fix up the children first, then decide what to do with the current node.

For each node, we recursively call pruneTree on the left and right children and reassign them: root.left = pruneTree(root.left). If a child's subtree was entirely zeros, the recursion returns null, which neatly detaches it.

After the children are cleaned, a node should be dropped only if it is now a leaf with value 0 — that is, root.val == 0 and both children are null. In that case we return null to our own parent; otherwise we keep the node.

Because we handle children before the parent, a deep all-zero branch collapses from the bottom up. Example: in [1, null, 0, 0, 1], the lone 0 leaf is dropped first, and the 0 above it keeps a 1 child so it survives, giving [1, null, 0, null, 1].

Each node is visited once, so pruning takes O(n) time.

Time: O(n) Space: O(h)