DEV Community

Wenqi Jiang
Wenqi Jiang

Posted on

Tree Pruning

Given a binary tree root, prune the tree so that subtrees containing all 0s are removed.

Constraints: n ≤ 100,000 where n is the number of nodes in root

Example 1
Input

root = [0, [1, null, null], [0, [1, [0, null, null], [0, null, null]], [0, null, null]]]
Output

[0, [1, null, null], [0, [1, null, null], null]]
Explanation
We do not remove the tree at the root or its right child because they still have a 1 as a descendant.

DFS:

  1. if current node is NULL, return NULL;
  2. if dfs(left) equals NULL, current left = NULL;
  3. if dfs(right) equals NULL, current right = NULL;
  4. if current node left & right both equal NULL, and current value is also 0, return NULL, else return current node
/**
 * class Tree {
 *     public:
 *         int val;
 *         Tree *left;
 *         Tree *right;
 * };
 */

Tree* solve(Tree* root) {
    if (root == NULL) {
        return NULL;
    }
    if (solve(root->left) == NULL) {
        root->left = NULL;
    }
    if (solve(root->right) == NULL) {
        root->right = NULL;
    }
    if (root->val == 0 && root->left == NULL && root->right == NULL) {
        return NULL;
    } else {
        return root;
    }
}
Enter fullscreen mode Exit fullscreen mode

Top comments (0)