Binary Trees — Delete Nodes and Return Forests

Today we are going to be looking at a binary tree problem.

The problem is asking us to find if any of our tree nodes have values that are also contained in our to_delete array that we are given. If we find these values we need to delete these nodes and return the various sub trees we are left with.

Alright so right away we know we are going to need to traverse our tree somehow and look at every node. Let’s create a dfs helper function.

First thing we always need to do is create a base case for ourselves so that our recursive function will stop and we will not be left in a dreaded infinite loop.

If we find a node that is equal to null then we will return null.

Next, we check our to_delete array and see if it contains this nodes value. If we find this value, then we know we want to delete this node, but we also know that we still need to check the nodes below this one and make sure to capture the branches, so we run our dfs function on both this current nodes .left and .right and give add the value of true, ensuring that in our else statement below, that those nodes get added to our results. We return, and then we set this nodes.left equal to null and this nodes.right equal to null. Lastly, we return our value of our node as null, successfully deleting it from our tree.

Now let’s look at the else if statement.

If our to_delete array does not contain this nodes.value then we are going set node.left equal to our dfs function, giving it our current nodes.left and false as parameters. Once we return, we set node.right equal to our dfs function as well, giving it our current nodes.right and false as parameters. This is the point from above we talked about. If add is true, (which is the second parameter we pass into our function) then we push this node into our results, and we return the current node.

Because it is so hard for me to visualize what this recursive function is actually doing, I like to write out an example to follow along and actually make sense of this.

Here is the final code: