【LeetCode-中等】814. 二叉树剪枝（C++实现）

题目链接

814. 二叉树剪枝

解题思路

对题目给定的树进行遍历。
遍历中，可以传递三个参数：当前节点、父亲节点、代表当前节点是父亲节点的左孩子还是右孩子（假设左孩子为 0，右孩子为 1）。

---

• 函数traversal() 对当前节点进行判断：

  • 为空则直接返回。
  • 不为空，则继续判断当前子树是否包含 1 ：
    • 若不包含 1 ，则根据参数 pos 确定是将父亲节点的哪一个子树剪掉。如果当前节点没有根节点（也就是当前节点就是题目所给的整棵树的根节点），那么将 root 设为 nullptr 后返回即可。
    • 如果包含 1 ，则继续处理当前节点的左、右子树。

• 对于判断当前子树是否包含 1 的函数 contain1()，根据题意，只需满足下面的三个条件之一即可：
  ① 当前节点的值为 1
  ② 当前节点的左子树中包含 1
  ③ 当前节点的右子树中包含 1

实现代码（C++）

class Solution {
public:
    bool contain1(TreeNode* root) {
        if (root == nullptr) {
            return false;
        }
        return root->val == 1 || contain1(root->left) || contain1(root->right);
    }
    void traversal(TreeNode* &root, TreeNode* parent, int pos) {
        if (root == nullptr) {
            return ;
        }
        if (!contain1(root)) { // 当前子树不包含1
            if (parent != nullptr) { // 不为整棵树的根节点
                if (pos == 0) { // 是左孩子
                    parent->left = nullptr;
                } else { // 是右孩子
                    parent->right = nullptr;
                }
            } else { // 是整棵树的根节点
                root = nullptr;
                return ;
            }
        } else {
            traversal(root->left, root, 0);
            traversal(root->right, root, 1);
        }
    }
    TreeNode* pruneTree(TreeNode* root) {
        traversal(root, nullptr, -1);
        
        return root;
    }
};