How to Prune a Binary Tree in C++?

We 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. (Recall that the subtree of a node X is X, plus every node that is a descendant of X.)

Example 1:

Binary Tree Pruning

Input: [1,null,0,0,1]
Output: [1,null,0,null,1]

Explanation:
Only the red nodes satisfy the property “every subtree not containing a 1”.
The diagram on the right represents the answer.

Example 2:

Binary Tree Pruning

Input: [1,0,1,0,0,0,1]
Output: [1,null,1,null,1]

Example 3:

Binary Tree Pruning

Input: [1,1,0,1,1,0,1,0]
Output: [1,1,0,1,1,null,1]

Binary Tree Pruning Algorithm in C++ Recursion

Like many other binary tree problems, pruning a binary tree can be solved elegantly using recursion. If the left sub tree can be pruned, we set the left pointer to NULL. If the right sub tree can be pruned, we can also set it to NULL.

If the current node value is 1, or, either of its children tree isn’t pruned, we return the current node, otherwise, this node can be pruned as well – set it to NULL.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    TreeNode* pruneTree(TreeNode* root) {
        if (root == nullptr) return nullptr;
        if (pruneTree(root->left) == nullptr) {
            root->left = nullptr;
        }
        if (pruneTree(root->right) == nullptr) {
            root->right = nullptr;
        }
        if ((root->val == 1) || (root->left != nullptr) || (root->right != nullptr)) {
            return root;
        }
        return nullptr;
    }
};

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    TreeNode* pruneTree(TreeNode* root) {
        if (root == nullptr) return nullptr;
        if (pruneTree(root->left) == nullptr) {
            root->left = nullptr;
        }
        if (pruneTree(root->right) == nullptr) {
            root->right = nullptr;
        }
        if ((root->val == 1) || (root->left != nullptr) || (root->right != nullptr)) {
            return root;
        }
        return nullptr;
    }
};

The complexity is O(N) where each node of the binary tree will be visited once. The space complexity is O(H) where H is the (maximum) height of the tree – as this is caused by the stack used in recursion.

Slightly different form, the above recursion can be coded as follows where the containsOne function will be explictly checking if the tree can be pruned.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    TreeNode* pruneTree(TreeNode* root) {
        return containsOne(root) ? root : nullptr;
    }
    
    bool containsOne(TreeNode* root) {
        if (root == nullptr) return false;
        bool left = containsOne(root->left);
        if (!left) {
            root->left = nullptr;
        }
        bool right = containsOne(root->right);
        if (!right) {
            root->right = nullptr;
        }
        return root->val == 1 || left || right;
    }
};

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    TreeNode* pruneTree(TreeNode* root) {
        return containsOne(root) ? root : nullptr;
    }
    
    bool containsOne(TreeNode* root) {
        if (root == nullptr) return false;
        bool left = containsOne(root->left);
        if (!left) {
            root->left = nullptr;
        }
        bool right = containsOne(root->right);
        if (!right) {
            root->right = nullptr;
        }
        return root->val == 1 || left || right;
    }
};

–EOF (The Ultimate Computing & Technology Blog) —