CODEWITHSUNDEEP
pythonbinary-tree
VVC
VVC

Reputation: 115

Removing a node from binary tree

I'm currently working on leetcode problem 366 where we have to find list of lists that contains values of leaves of each generation. I wanted to achieve this by recursion where if a node does not have left or right child, the value is recorded then the node removed by setting it to None. Here is my code:

def findLeaves(self, root: Optional[TreeNode]) -> List[List[int]]:
    leaf_list = []
    sub_list = []
    def traverse(node):
        if node == None:
            return
        if node.left == None and node.right == None:
            sub_list.append(node.val)
            node = None
            return 
        traverse(node.left)
        traverse(node.right)
        return root
    
    while True:
        if root == None:
            break
        sub_list = []
        traverse(root)
        leaf_list.append(sub_list)
        print(leaf_list)
    return leaf_list

The problem seems to be that when a certain node is set to None, that change isn't retained. Why is it that I can't set a node to None to remove it?

Thanks

Upvotes: 0

Views: 540

Answers (1)

trincot
trincot

Reputation: 350270

The tree can only be mutated when you assign to one if its node's attributes. An assignment to a variable, only changes what the variable represents. Such assignment never impacts whatever previous value that variable had. Assigning to a variable is like switching from one value to another without affecting any data structure. So you need to adapt your code such that the assignment of None is done to a left or right attribute.

The exception is for the root node itself. When the root is a leaf, then there is no parent to mutate. You will then just discard the tree and switch to an empty one (None).

One way to achieve this, is to use the return value of traverse to update the child-reference (left or right) that the caller of traverse needs to update.

Here is your code with those adaptations:

def findLeaves(root):
    sub_list = []
    def traverse(node):
        if not node:
            return
        if not node.left and not node.right:
            sub_list.append(node.val)
            return  # By returning None, the parent can remove it
        node.left = traverse(node.left)  # Assign the returned node reference
        node.right = traverse(node.right)
        return node  # Return the node (parent does not need to remove it)
    
    leaf_list = []
    while root:
        sub_list = []
        root = traverse(root)
        leaf_list.append(sub_list)
    return leaf_list

Upvotes: 1

Related Questions