Tree Pruning Algorithm in Swift

Hello passionate coders, Leo here. The topic that we will study today is the binary Tree Pruning Algorithm in Swift.

We’ll continue the Algorithms series here. The algorithm involves binary tree data structure and some recursion, stay here to explore this challenge. The problem statement you can find and test at the awesome Leetcode site.

This kind of challenge is important for every developer who wants the open his mind to new sets of problems, ultimately having problem-solving skills. And why does that matter? Because problem-solving is what we do every day on the job. This way practice solving different set of problems with different approaches makes your life easier. And it’s fun after all.

The painting chosen today is *”Pastoral, Land Girls Pruning at East Malling”* by Evelyn Dunbar in 1944. She was a British artist, illustrator, and teacher. And she is notable for recording women’s contributions to World War II on the United Kingdom home front, particularly the work of the Women’s Land Army. One of that amazing woman’s contributions was… tree pruning. Got it?

Let’s code!

The problem – The Tree Pruning Algorithm in Swift

Image that you are given the head node root of a binary tree, where additionally every node’s value is either a 0 or a 1. Return the same tree where every subtree (of the given tree) not containing a 1 has been removed.

To better understand the problem here are some visualizations of what we have to do:

and

You can check the images above if you are having trouble understanding the target of our challenge.

Approaching the problem

When you are trying to solve any kind of problem it’s very important you pay attention to the hints that the problem gives to you. The raw Swift function that Leetcode gives to you is already a hint that you can use to solve the problem. Let me explain after you read the code above:

class Solution {
   func pruneTree(_ root: TreeNode?) -> TreeNode? {

   }
}

Here you can see clearly a hint about how this could be solved. If the function parameter is a TreeNode? and the return type is a TreeNode? this is a very strong hint pointing to a recursive solution. Another thing that we could think about the solution is that binary tree problem solutions regularly use recursion.

With that in mind, the first approach I thought of was to do a recursive method to check every sub from root to the leaf, and checking if it was a “prunable” node I could just nil it. This was getting bigger than I expected and I was feeling that was a dead end.

But what if, I invert the check? Instead of checking from the root of each subtree, I could check from the leaves and go up each tree level. Yes! That should work. So let’s write down the solution to the code below:

guard let root = root else { return nil } // mark 1

root.left = pruneTree(root.left) // mark 2
root.right = pruneTree(root.right) // mark 2

if root.val == 0 && root.left == nil && root.right == nil { mark 3
    return nil
}

return root // mark 4
}

The explanation is simple:

1. Mark 1 – This is a handy guard to don’t need to work with optional all around the place.
2. Mark 2 – Recursively traverse the binary to the leaf nodes.
3. Mark 3 – Check if the node has value == 0, if yes check if the node has nil left and right nodes. If both children are nil you can prune this node (return nil)
4. Mark 4 -If the condition above fails that means you cannot prune the node, so you return it.

That’s it! A very straightforward algorithm, right?

Algorithms: Continue Studying

Aside from trees, one data structure you should really master is arrays. Techniques like two pointers are your bread and butter of arrays problems, and you can practice that in the article squares of a sorted array problem. There you will study how important is to know the two-pointers technique even for the easiest problems.

Working with graphs in interviews is hard. It is no news that in an interview, especially on Google, you could be asked how to solve a complex graph algorithm under a very short time limit. To practice that I an article about the A* graph algorithm and how to solve a children’s puzzle with it.

Summary – Tree Pruning Algorithm in Swift

Today we explore a simple solution of a recursive binary tree pruning in Swift. I’m excited to hear about what was your solution in this case.

That’s all my people, I hope you liked reading this article as much as I enjoyed writing it. If you want to support this blog you can Buy Me a Coffee or leave a comment saying hello. You can also sponsor posts and I’m open to freelance writing! You can reach me on LinkedIn or Twitter and send me an e-mail through the contact page.

Thanks for reading and… That’s all folks.

Credits: image image of the trees