814. Binary Tree Pruning

Problem Description

In this problem, we are given the root of a binary tree. Our task is to modify the tree by removing all subtrees that do not have at least one node with the value 1. A subtree, as defined in the problem, is composed of a node and all of its descendants. The goal is to return the same tree, but with certain nodes pruned so that any remaining subtree contains the value 1.

Flowchart Walkthrough

First, let's determine the appropriate algorithm using the Flowchart. Here’s a step-by-step breakdown:

Is it a graph?

Yes: The problem deals with a binary tree, which is a specific type of graph.

Is it a tree?

Yes: Specifically, the binary tree is given, making it a tree structure.

Is the problem related to directed acyclic graphs (DAGs)?

No: While trees can be considered a type of DAG, the problem is specifically about pruning a binary tree based on certain node conditions, not about properties or problems typical to generic DAGs like topological sorting.

Is the problem related to shortest paths?

No: The task does not involve finding shortest paths; it revolves around modifying the tree structure based on node values.

Conclusion: Following the flowchart, the problem being directly related to tree operations and not fitting the criteria for DAG-related algorithms or shortest path problems points toward using Depth-First Search (DFS). DFS is suitable for tree operations as it allows exploring all nodes and making decisions (prune or retain) based on child nodes recursively, which aligns with the problem's requirements of pruning the tree based on child values.

Intuition

The intuition behind the solution is based on a post-order traversal of the tree, where we visit the child nodes before dealing with the parent node. The recursive function pruneTree works bottom-up:

If the current node is None, there is nothing to do, so we return None.
Recursively call pruneTree on the left child of the current node. The return value will be the pruned left subtree or None if the left subtree doesn't contain a 1.
Perform the same operation for the right child.
Once both left and right subtrees are processed, check if the current node's value is 0 and that it has no left or right children that contain a 1 (as they would have been pruned if that was the case).
If the current node is a 0 node with no children containing a 1, it is pruned as well by returning None.
If the current node should remain, we return the current node with its potentially pruned left and right subtrees.

This recursive approach ensures that we remove all subtrees which do not contain a 1 by considering each node's value and the result of pruning its subtrees.

Learn more about Tree, Depth-First Search and Binary Tree patterns.

Solution Approach

The implementation of the solution employs the following concepts:

Recursion: The core mechanism for traversing and modifying the tree is the recursive function pruneTree.
Post-Order Traversal: This pattern is used where we process the children nodes before the parent node, which is essential in this problem because the decision to prune a parent node depends on the status of its children.
Tree Pruning: Nodes are removed from the tree based on a certain condition (in this case, the absence of the value 1 in the subtree rooted at that node).

Here’s a step-by-step walkthrough of how the algorithm works using the provided Python code:

Recursively call pruneTree(root) with the original root of the tree. If the root is None, it returns None, signifying an empty tree.
The pruneTree function is called on root.left, pruning the left subtree. The assignment root.left = self.pruneTree(root.left) ensures that the current node's left pointer is updated to the result of pruning its left subtree - either a pruned subtree or None.
The same operation is performed on root.right to prune the right subtree.
After both subtrees are potentially pruned, we check the current node. If root.val is 0, and both root.left and root.right are None (meaning they were pruned or originally empty), this node must also be pruned. Therefore, the function returns None.
If the current node is not pruned (i.e., it has the value 1 or it has a non-pruned subtree), it is returned as is.

The result of the recursive function is a reference to the root of the pruned tree since each recursive call potentially modifies the left and right children of the nodes by returning either a pruned subtree or None.

The use of recursion here allows the solution to smoothly handle the nested/tree-like structure of the problem and prune nodes appropriately based on the state of their descendants.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through a small example to illustrate the solution approach. Consider a binary tree represented by the following structure:

We need to prune this tree by removing all subtrees that do not have at least one node with the value 1. We will use a post-order traversal method to achieve this.

We start at the leftmost node, which has the value 0. Since it has no children, it does not contain a 1, so we prune this node, and it becomes None.
Moving up to its parent, which also has a value of 0. It has no right child and its left child was pruned in step 1. As there is no 1 in the subtree, we prune this node as well.
Now, we visit the left child of the root, which again has a value of 0. Its left child (subtree from steps 1 and 2) has already been pruned. This node does not have a right child. Therefore, we prune this node, and the entire left subtree of the root is now None.
We move to the right subtree of the root. The left child of the root's right child has the value 0. This node has a single child with the value 1, so we keep this subtree intact.
Moving to the right child of the root's right child, it has the value 1, so we keep this node.
Looking at the parent of the nodes from steps 4 and 5, which is the right child of the root with the value 1, since it contains 1 and its subtrees (which are its children nodes) also contain the value 1 or are None, we keep this subtree intact.
Finally, we move to the root of the tree. Since its right subtree contains a 1, we keep the right subtree. The left subtree was pruned earlier. The root itself has a value 0, but because it has a non-empty right subtree that contains a 1, we do not prune the root.

After the pruning operation, our tree looks like this:

The pruned tree now only contains subtrees which have at least one node with the value 1 as required by the problem statement.

Solution Implementation

1# Definition for a binary tree node.
2class TreeNode:
3    def __init__(self, val=0, left=None, right=None):
4        self.val = val
5        self.left = left
6        self.right = right
7
8class Solution:
9    def prune_tree(self, root: TreeNode) -> TreeNode:
10        # If the current node is None, no pruning needed, return None.
11        if root is None:
12            return None
13
14        # Prune the left subtree and assign the result to the left child.
15        root.left = self.prune_tree(root.left)
16      
17        # Prune the right subtree and assign the result to the right child.
18        root.right = self.prune_tree(root.right)
19      
20        # Check if the current node is a leaf with value 0, prune it by returning None.
21        if root.val == 0 and root.left is None and root.right is None:
22            return None
23
24        # Return the current node as it is valid in the pruned tree.
25        return root
26
27    # Alias the method to the originally provided method name to comply with LeetCode interface.
28    pruneTree = prune_tree
29

1/**
2 * Definition for a binary tree node.
3 * public class TreeNode {
4 *     int value; // TreeNode's value
5 *     TreeNode left; // left child
6 *     TreeNode right; // right child
7 *     TreeNode() {} // default constructor
8 *     TreeNode(int value) { this.value = value; } // constructor with value
9 *     // constructor with value, left child, and right child
10 *     TreeNode(int value, TreeNode left, TreeNode right) {
11 *         this.value = value;
12 *         this.left = left;
13 *         this.right = right;
14 *     }
15 * }
16 */
17class Solution {
18    /**
19     * Prunes a binary tree by removing all subtrees that do not contain the value 1.
20     * A subtree rooted at node containing 0s only is pruned.
21     *
22     * @param root The root of the binary tree.
23     * @return The pruned binary tree's root.
24     */
25    public TreeNode pruneTree(TreeNode root) {
26        // Base case: if the node is null, return null (i.e., prune the null subtree)
27        if (root == null) {
28            return null;
29        }
30      
31        // Recursively prune the left subtree
32        root.left = pruneTree(root.left);
33      
34        // Recursively prune the right subtree
35        root.right = pruneTree(root.right);
36      
37        // If the current node's value is zero and it doesn't have any children,
38        // it is a leaf node with value 0, thus should be pruned (returned as null)
39        if (root.value == 0 && root.left == null && root.right == null) {
40            return null;
41        }
42      
43        // Otherwise, return the current (possibly pruned) node.
44        return root;
45    }
46}
47

1// Definition for a binary tree node.
2struct TreeNode {
3    int val; // Value of the node
4    TreeNode *left; // Pointer to the left child
5    TreeNode *right; // Pointer to the right child
6
7    // Constructor initializes the node with a default value and null child pointers
8    TreeNode() : val(0), left(nullptr), right(nullptr) {}
9
10    // Constructor initializes the node with a given value and null child pointers
11    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
12
13    // Constructor initializes the node with a given value and given left and right children
14    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
15};
16
17class Solution {
18public:
19    // Function to prune a binary tree. If a subtree does not contain the value '1'
20    // it will remove that subtree.
21    TreeNode* pruneTree(TreeNode* root) {
22        // If the current node is null, return null (base case)
23        if (!root) return nullptr;
24
25        // Recursively prune the left subtree
26        root->left = pruneTree(root->left);
27
28        // Recursively prune the right subtree
29        root->right = pruneTree(root->right);
30
31        // If the current node's value is 0 and both left and right subtrees are null (pruned away),
32        // then remove this node as well by returning null
33        if (root->val == 0 && !root->left && !root->right) return nullptr;
34
35        // If the current node is not pruned, return the node
36        return root;
37    }
38};
39

1// Definition for a binary tree node.
2interface TreeNode {
3    val: number;
4    left: TreeNode | null;
5    right: TreeNode | null;
6}
7
8/**
9 * Recursively prunes a binary tree, removing all subtrees that contain only 0s.
10 * @param {TreeNode | null} node - The root node of the binary tree or subtree being pruned.
11 * @returns {TreeNode | null} - The pruned tree's root node, or null if the entire tree is pruned.
12 */
13function pruneTree(node: TreeNode | null): TreeNode | null {
14    // Base case: if the current node is null, just return it.
15    if (node === null) {
16        return node;
17    }
18
19    // Recursively prune the left and right subtrees.
20    node.left = pruneTree(node.left);
21    node.right = pruneTree(node.right);
22
23    // If the current node's value is 0 and it has no children, prune it by returning null.
24    if (node.val === 0 && node.left === null && node.right === null) {
25        return null;
26    }
27
28    // If the current node has a value of 1 or has any children, return the node itself.
29    return node;
30}
31

Time and Space Complexity

The time complexity of the code is O(n), where n is the number of nodes in the binary tree. This is because the algorithm must visit each node exactly once to determine whether it should be pruned.

The space complexity of the code is O(h), where h is the height of the binary tree. This represents the space taken up by the call stack due to recursion. In the worst-case scenario, the function could be called recursively for each level of the tree, resulting in a stack depth equal to the height of the tree. For a balanced tree, this would be O(log n), but for a skewed tree, it could be O(n).

Learn more about how to find time and space complexity quickly using problem constraints.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan: