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:
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:
Input: [1,0,1,0,0,0,1]
Output: [1,null,1,null,1]Example 3:
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) —
loading...
Last Post: Algorithm to Compute the Shortest Distance between Points on Two Lines
Next Post: Understanding The Math Before Investing In Stocks