Removing a node from binary tree

Question

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

• Scott Hunter: There is a difference between changing node and changing the part of the tree that points to node.
• VVC: could you please elaborate on that? If at 3rd recursion node is root.left.left.left and that node is an instance of class TreeNode, wouldn't setting node = None equate removing root.left.left.left?
• Scott Hunter: Obviously not, or you wouldn't be here. If you want root.left.left.left to be None, you have to assign None to that.

Answer 1

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