Leetcode 814. Binary Tree Pruning

Problem Explanation

The problem is about pruning a binary tree. This means that you have to return a modified version of the provided binary tree, where every subtree that does not contain a 1 is removed.

A subtree of a node X is defined as X, plus every node that is a descendant of X. The root node of a binary tree is always considered to be a subtree itself. The binary tree will have at most 100 nodes, and the value of each node will only be 0 or 1.

This is a recursion problem where we need to traverse the tree, prune out subtrees that do not contain a 1, and then return the pruned tree.

Walkthrough of an example: For example, given a binary tree with root node structure [1, 0, 1, 0, 0, 0, 1], the output should be [1, null, 1, null, 1] after pruning.

Approach to the Solution

To prune the binary tree, we can use Depth-First Search (DFS) recursion to traverse both left and right subtrees. During the traversal, if we reach a leaf node which does not contain 1, we prune it out by returning null to its parent node. Finally, we return the root of the binary tree.

Solution

Python

1
2python
3class Solution:
4    def pruneTree(self, root):
5        if root is None:
6            return None
7        
8        # Recursively prune the left and right subtrees
9        root.left = self.pruneTree(root.left)
10        root.right = self.pruneTree(root.right)
11        
12        # If both left and right subtrees don't contain a 1 and the current node value is 0
13        # we prune it by returning null to its parent node
14        if root.left is None and root.right is None and root.val == 0:
15            return None
16        
17        return root

Java

1
2java
3class Solution {
4    public TreeNode pruneTree(TreeNode root) {
5        if (root == null)
6            return null;
7        
8        // Recursively prune the left and right subtrees
9        root.left = pruneTree(root.left);
10        root.right = pruneTree(root.right);
11        
12        // If both left and right subtrees don't contain a 1 and the current node value is 0
13        // we prune it by returning null to its parent node
14        if (root.left == null && root.right == null && root.val == 0)
15            return null;
16        
17        return root;
18    }
19}

JavaScript

1
2javascript
3class Solution {
4    pruneTree(root) {
5        if (root === null)
6            return null;
7        
8        // Recursively prune the left and right subtrees
9        root.left = this.pruneTree(root.left);
10        root.right = this.pruneTree(root.right);
11        
12        // If both left and right subtrees don't contain a 1 and the current node value is 0
13        // we prune it by returning null to its parent node
14        if (root.left === null && root.right === null && root.val === 0)
15            return null;
16        
17        return root;
18    }
19}

C++

1
2c++
3class Solution {
4public:
5    TreeNode* pruneTree(TreeNode* root) {
6        if (root == nullptr)
7            return nullptr;
8        
9        // Recursively prune the left and right subtrees
10        root->left = pruneTree(root->left);
11        root->right = pruneTree(root->right);
12        
13        // If both left and right subtrees don't contain a 1 and the current node value is 0
14        // we prune it by returning null to its parent node
15        if (root->left == nullptr && root->right == nullptr && root->val == 0)
16            return nullptr;
17        
18        return root;
19    }
20};

C#

1
2csharp
3public class Solution {
4    public TreeNode PruneTree(TreeNode root) {
5        if (root == null)
6            return null;
7        
8        // Recursively prune the left and right subtrees
9        root.left = PruneTree(root.left);
10        root.right = PruneTree(root.right);
11        
12        // If both left and right subtrees don't contain a 1 and the current node value is 0
13        // we prune it by returning null to its parent node
14        if (root.left == null && root.right == null && root.val == 0)
15            return null;
16        
17        return root;
18    }
19}

Time and Space Complexity Analysis

In our solution, we are traversing the binary tree once and in each traversal, we process each node only once. Therefore, the time complexity of the solution is O(n), where n is the number of nodes in the binary tree.

As for the space complexity, our solution uses recursion which causes the function call stack to store a maximum number of function calls equal to the height of the tree in the worst case. Therefore, the space complexity is O(h), where h is the height of the binary tree.

These solutions are efficient and handle the problem quite well. To further optimize the solution, one could consider using a stack data structure instead of recursion for DFS traversal, which could reduce the maximum space usage.

These are some of the best practices to follow while solving similar tree pruning problems:

  1. Always check if the root of the tree or subtree is null. This can help avoid null pointer exceptions.

  2. Try to use recursion where possible. It reduces the complexity of the problem and makes the solution easier to understand and implement.

  3. While using DFS traversal, always take care of the order in which you visit the nodes. It will decide how the pruning takes place.

  4. Always test your solution with different inputs to ensure it works for all edge cases. For example, consider testing your solution with a binary tree that already doesn't have any 0 nodes or a binary tree that has all 0 nodes.

These tips will help you write robust solutions for binary tree pruning problems and similar challenges.


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫