Binary Tree Pruning

Med

#0768Time: O(N), where N is the number of nodes. The tree is traversed twice, so the complexity is O(N) + O(N) = O(N).Space: O(N) in the worst case. This is dominated by the `HashMap` which stores an entry for each node. The recursion stack also contributes O(H) space, where H is the tree height, which can be up to N for a skewed tree.1 company

Algorithms

Depth-First Search

Data structures

Tree Binary Tree

Companies

Hulu

Prompt

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.

A subtree of a node node is node plus every node that is a descendant of node.

Example 1:

Input: root = [1,null,0,0,1]
Output: [1,null,0,null,1]
Explanation: 
Only the red nodes satisfy the property "every subtree not containing a 1".
The diagram on the right represents the answer.

Example 2:

Input: root = [1,0,1,0,0,0,1]
Output: [1,null,1,null,1]

Example 3:

Input: root = [1,1,0,1,1,0,1,0]
Output: [1,1,0,1,1,null,1]

Constraints:

The number of nodes in the tree is in the range [1, 200].
Node.val is either 0 or 1.

Approaches

2 approaches with complexity analysis and trade-offs.

This approach separates the problem into two distinct steps. First, it traverses the tree to identify and mark all subtrees that contain a '1'. Second, it performs another traversal to prune the subtrees that were not marked.

Algorithm

1. Create a HashMap<TreeNode, Boolean> to memoize whether a subtree contains a 1.
1. Perform a first pass using a post-order traversal (hasOne function) to populate the map. For each node, the value stored is true if node.val == 1 or if either of its children's subtrees contains a 1.
1. Check the map for the root. If the entire tree does not contain a 1, return null.
1. Perform a second pass (e.g., pre-order traversal, pruneNodes function) to modify the tree. For each node, check its children against the map. If a child's subtree is marked as not containing a 1, set the child pointer to null.
1. Return the modified root.

Walkthrough

In this method, we first need to figure out for each node whether the subtree rooted at that node contains a 1. A post-order traversal is suitable for this, as a parent's status depends on its children's status. We can store this boolean information in a HashMap with nodes as keys.

Once the map is populated, we perform a second traversal (e.g., pre-order). For each node, we check the map for its children. If a child's subtree is marked as not containing a '1', we sever the link to that child, effectively pruning it.

Finally, the root itself might need to be pruned if its entire tree contains no '1's, which we can check from the map.

Here is the Java implementation:

1/**2 * Definition for a binary tree node.3 * public class TreeNode {4 *     int val;5 *     TreeNode left;6 *     TreeNode right;7 *     TreeNode() {}8 *     TreeNode(int val) { this.val = val; }9 *     TreeNode(int val, TreeNode left, TreeNode right) {10 *         this.val = val;11 *         this.left = left;12 *         this.right = right;13 *     }14 * }15 */16class Solution {17    // A map to store whether a subtree at a node contains a 1.18    private Map<TreeNode, Boolean> containsOneMap = new HashMap<>();19 20    public TreeNode pruneTree(TreeNode root) {21        // First pass: check which subtrees contain a 1.22        boolean rootHasOne = hasOne(root);23        24        // If the entire tree has no 1s, return null.25        if (!rootHasOne) {26            return null;27        }28        29        // Second pass: prune the tree based on the map.30        return pruneNodes(root);31    }32 33    // Post-order traversal to populate the map.34    private boolean hasOne(TreeNode node) {35        if (node == null) {36            return false;37        }38        boolean leftHasOne = hasOne(node.left);39        boolean rightHasOne = hasOne(node.right);40        boolean result = node.val == 1 || leftHasOne || rightHasOne;41        containsOneMap.put(node, result);42        return result;43    }44 45    // Pre-order traversal to prune nodes.46    private TreeNode pruneNodes(TreeNode node) {47        if (node == null) {48            return null;49        }50        // Check children and prune if their subtrees have no 1s.51        if (node.left != null && !containsOneMap.get(node.left)) {52            node.left = null;53        }54        if (node.right != null && !containsOneMap.get(node.right)) {55            node.right = null;56        }57        58        // Recurse on the remaining children.59        pruneNodes(node.left);60        pruneNodes(node.right);61        62        return node;63    }64}

Complexity

Time

O(N), where N is the number of nodes. The tree is traversed twice, so the complexity is O(N) + O(N) = O(N).

Space

O(N) in the worst case. This is dominated by the `HashMap` which stores an entry for each node. The recursion stack also contributes O(H) space, where H is the tree height, which can be up to N for a skewed tree.

Trade-offs

Pros

The logic is separated into two distinct, easy-to-understand phases: analysis and modification.
Can be easier to reason about for those less comfortable with complex recursive return values.

Cons

Requires two full traversals of the tree, making it less efficient in terms of constant factors.
Uses O(N) extra space for the hash map, which is suboptimal.

This highly efficient approach uses a single post-order traversal to solve the problem. The decision to prune a node is made immediately after processing its children. The recursive function modifies the tree in-place and returns the potentially pruned subtree to its parent.

Algorithm

1. The solution is a recursive function that takes a node and returns the root of the pruned subtree.
1. Base Case: If the current node is null, return null.
1. Recursive Step (Post-order): Recursively call the function on the left and right children and update the current node's left and right pointers with the results. node.left = pruneTree(node.left) and node.right = pruneTree(node.right).
1. Pruning Condition: After the recursive calls, check if the current node should be pruned. The condition for pruning is: node.left == null, node.right == null, and node.val == 0.
1. Return Value: If the pruning condition is met, return null. Otherwise, return the current node.

Walkthrough

The core idea is that a subtree should be removed if and only if it does not contain a 1. This property can be checked recursively. A post-order traversal is a natural fit because to decide about a parent node, we must first have information about its children.

The recursive function, say pruneTree(node), will do the following:

Recursively call itself for the left and right children. This will prune the left and right subtrees first.
The calls pruneTree(node.left) and pruneTree(node.right) will return null if the respective subtrees are completely pruned, or the root of the pruned subtree otherwise. We update node.left and node.right with these results.
After the children have been processed, we look at the current node. If its value is 0 and both of its children are now null (either originally or after pruning), then this node itself is the root of a subtree with no 1s. In this case, the function should return null to its caller, effectively pruning this node.
Otherwise, the node is part of a valid subtree and should be kept. The function returns the node itself.

This single pass elegantly combines checking and pruning.

Here is the concise Java implementation:

1/**2 * Definition for a binary tree node.3 * public class TreeNode {4 *     int val;5 *     TreeNode left;6 *     TreeNode right;7 *     TreeNode() {}8 *     TreeNode(int val) { this.val = val; }9 *     TreeNode(int val, TreeNode left, TreeNode right) {10 *         this.val = val;11 *         this.left = left;12 *         this.right = right;13 *     }14 * }15 */16class Solution {17    public TreeNode pruneTree(TreeNode root) {18        // Base case for recursion19        if (root == null) {20            return null;21        }22        23        // Process children first (post-order traversal)24        root.left = pruneTree(root.left);25        root.right = pruneTree(root.right);26        27        // Check if the current node should be pruned28        // A node is pruned if it's a leaf and its value is 0.29        // A node can become a leaf after its children are pruned.30        if (root.left == null && root.right == null && root.val == 0) {31            return null;32        }33        34        // Otherwise, keep the node35        return root;36    }37}

Complexity

Time

O(N), where N is the number of nodes in the tree. Each node is visited exactly once.

Space

O(H), where H is the height of the tree. This space is used by the recursion call stack. In the worst case of a completely unbalanced tree, the height H can be equal to N, leading to O(N) space complexity.

Trade-offs

Pros

Extremely efficient, solving the problem in a single pass.
Space-efficient, using only O(H) space for the recursion stack instead of O(N) auxiliary space.
The code is very concise and elegant.

Cons

For extremely deep trees (not an issue given the problem constraints), recursion could lead to a stack overflow.
The logic of returning a node or null to modify the parent's link might be slightly less direct to grasp than a two-pass approach.

Solutions

Solution

1/** * Definition for a binary tree node. * public class TreeNode { * int val; * TreeNode left; * TreeNode right; * TreeNode() {} * TreeNode(int val) { this.val = val; } * TreeNode(int val, TreeNode left, TreeNode right) { * this.val = val; * this.left = left; * this.right = right; * } * } */ 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 ; } }

1/** * Definition for a binary tree node. * function TreeNode(val, left, right) { * this.val = (val===undefined ? 0 : val) * this.left = (left===undefined ? null : left) * this.right = (right===undefined ? null : right) * } */ /** * @param {TreeNode} root * @return {TreeNode} */ var pruneTree =2  function (root) {3    if (!root) return null;4    root.left = pruneTree(root.left);5    root.right = pruneTree(root.right);6    if (root.val == 0 && !root.left && !root.right) {7      return null;8    }9    return root;10  };

1/** * Definition for a binary tree node. * struct TreeNode { * int val; * TreeNode *left; * TreeNode *right; * TreeNode() : val(0), left(nullptr), right(nullptr) {} * TreeNode(int x) : val(x), left(nullptr), right(nullptr) {} * TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {} * }; */ class Solution { public: TreeNode * pruneTree ( TreeNode * root ) { if ( ! root ) return nullptr ; root -> left = pruneTree ( root -> left ); root -> right = pruneTree ( root -> right ); if ( ! root -> val && ! root -> left && ! root -> right ) return nullptr ; return root ; } };

1# Definition for a binary tree node. # class TreeNode: # def __init__(self, val=0, left=None, right=None): # self.val = val # self.left = left # self.right = right class Solution : def pruneTree ( self , root : Optional [ TreeNode ]) -> Optional [ TreeNode ]: if root is None : return None root . left = self . pruneTree ( root . left ) root . right = self . pruneTree ( root . right ) if root . val == 0 and root . left is None and root . right is None : return None return root

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem