Issues with Binary Tree Pruning implementation in Python

Question

I'm currently working on LeetCode problem 814, "Binary Tree Pruning," and I've encountered some difficulties with my implementation. The problem requires returning a modified binary tree where every subtree not containing a value of 1 has been removed.

Here's my initial code:

# 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

        self.pruneTree(root.left)
        self.pruneTree(root.right)

        if root.val == 0 and root.left is None and root.right is None:
            root = None
            return 

        return root

Unfortunately, my implementation did not produce the expected results and caused errors. However, when I consulted with others, they suggested the following code, which works correctly:

# 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:
            root = None
            return 

        return root

I'm having trouble understanding why my initial solution, which doesn't utilize the root.left and root.right variables, didn't work as expected, while the modified solution provided by others does. I would appreciate it if someone could shed some light on this and help me understand the underlying issue.

Thank you in advance for your assistance!

Publishing content generated by ChatGPT (not only answers) [isn't allowed currently](https://meta.stackoverflow.com/questions/421831/temporary-policy-chatgpt-is-banned) — Michael Butscher, Jul 19 '23 at 14:16
In your code, you weren't modifying the branches - the `return` value was not used at all, it was just discarded. With new one, you actually say "in the current object, I want to change left to whatever recursion on left gives me and right to whatever recursion on right gives me". Also, the root = None does nothing, the point is the `return` without anything - which is equal to `return None`. `root` is just local variable there, only the returned value matters — h4z3, Jul 19 '23 at 14:18

score 1 · Accepted Answer · answered Jul 19 '23 at 14:16

you perform the recursive calls to self.pruneTree(root.left) and self.pruneTree(root.right), but you don't assign the return values back to the root.left and root.right!

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)

The key difference is that the version above keeps track of the pruned subtrees by updating the left and right child pointers of the current node.

score 0 · Answer 2 · answered Jul 19 '23 at 15:11

As your code never assigns anything to an attribute, we can already be sure the tree will not change.

The assignment root = None will only set the value of a local name. It doesn't affect the tree; it doesn't change a left or right attribute. So doing:

    root = None
    return

Has no meaningful difference with just doing:

    return

About this aspect, see also How do I pass a variable by reference.

Issues with Binary Tree Pruning implementation in Python

2 Answers2