【面试题 & LeetCode 814】Binary Tree Pruning

题目描述

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.)

思路

dfs

代码

代码一：

/**
 * 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 == NULL) return root;
        root->left = pruneTree(root->left);
        root->right = pruneTree(root->right);
        if (!root->left && !root->right && root->val == 0) return NULL;
        return root;
    }
};

代码二：

/**
 * 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:
    map<TreeNode*, int> mp;
    TreeNode* pruneTree(TreeNode* root) {
        if (root == NULL) return root;
        if (has1(root->left) == false) root->left = NULL;
        if (has1(root->right) == false) root->right = NULL;
        pruneTree(root->left);
        pruneTree(root->right);
        return root;
    }
    
    bool has1(TreeNode* root) {
        if (mp.count(root) != 0) return mp[root];
        if (root == NULL) return false;
        if (root->val == 1) return true;
        return mp[root] = (has1(root->left) || has1(root->right));
    }
};