Tree Pruning algorithm in Swift

Subscribe to my newsletter and never miss my upcoming articles

Hello passionate coders, Leo here.

Today we'll continue the Algorithms series here. The algorithm envolve 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 want the open his mind to new sets of problems, ultimately having a problem solving skills. And why that matters? Because problem solving is what we do everyday in the job. This way practice solving different set of problems with different approaches makes you life easier. And it's fun after all.

The painting chosen today is "Pastoral, Land Girls Pruning at East Malling" from 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 contributions was... tree pruning. Got it?

Let's code!

The problem

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 some visualizations of what we have to do:

Screen Shot 2021-04-28 at 07.53.13.png

Screen Shot 2021-04-28 at 07.54.29.png

You can check the images above if you are having troubles 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 on the hints that the problem give to you. The raw Swift function that Leetcode give 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? too this is a very strong hint pointing to a recursive solution . Another think 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 was 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 checking from the root each subtree, I could check from the leafs and going up each tree level. Yes! That should work. So let's write down the solution to 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 fail that means you cannot prune the node, so you return it.

That's it! Very straightforward algorithm right?


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

That's all my people, I hope you liked this as I liked writing. If you want to support this blog you can Buy Me a Coffee or just leave a comment saying hello.

Thanks for the reading and... That's all folks.

credits: image of the trees cover photo

No Comments Yet