【东华大学oj】二叉树剪枝

二叉树剪枝

时间限制: 1s

类别: DS：树->力扣

问题描述

给你二叉树的根结点 root ，树的每个结点的值要么是 0 ，要么是 1 。

返回移除了所有不包含 1 的子树的二叉树。

节点 node 的子树为 node 本身加上所有 node 的后代。

示例 1：

输入：

5

1 null 0 0 1

输出：

1 null 0 null 1

解释：

只有红色节点满足条件“所有不包含 1 的子树”。 右图为返回的答案。

示例 2：

输入：

7

1 0 1 0 0 0 1

输出：

1 null 1 null 1

可以使用以下函数输入一棵树：

struct TreeNode

{

   int val;

   TreeNode *left;

   TreeNode *right;

   TreeNode() : val(0), left(NULL), right(NULL){}

   TreeNode(int x) : val(x), left(NULL), right(NULL) {}

   TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}

};

TreeNode* inputTree()

{

   int n,count=0;

   char item[100];

   cin>>n;

   if (n==0)

       return NULL;

   cin>>item;

   TreeNode* root = new TreeNode(atoi(item));

   count++;

   queue<TreeNode*> nodeQueue;

   nodeQueue.push(root);

   while (count<n)

   {

       TreeNode* node = nodeQueue.front();

       nodeQueue.pop();

       cin>>item;

       count++;

       if (strcmp(item,"null")!=0)

       {

           int leftNumber = atoi(item);

           node->left = new TreeNode(leftNumber);

           nodeQueue.push(node->left);

       }

       if (count==n)

           break;

       cin>>item;

       count++;

       if (strcmp(item,"null")!=0)

       {

           int rightNumber = atoi(item);

           node->right = new TreeNode(rightNumber);

           nodeQueue.push(node->right);

       }

   }

   return root;

}

输入说明

首先输入结点的数目n（注意，这里的结点包括题中的null空结点）

然后输入n个结点的数据，需要填充为空的结点，输入null。

注意，输入结点的数据只有0 或 1。

提示：

   树中节点的数目在范围 [1, 200] 内

   Node.val 为 0 或 1

输出说明

输出结果，需要填充为空的结点，输出null，每个数据的后面跟一个空格。

#include <iostream>
#include <cstring>
#include <queue>
using namespace std;

// Definition for a binary tree node.
struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode() : val(0), left(NULL), right(NULL){}
    TreeNode(int x) : val(x), left(NULL), right(NULL){}
    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right){}
};

TreeNode* inputTree() {
    int n, count = 0;
    char item[100];
    cin >> n;
    if (n == 0)
        return NULL;
    cin >> item;
    TreeNode* root = new TreeNode(atoi(item));
    count++;
    queue<TreeNode*> nodeQueue;
    nodeQueue.push(root);
    while (count < n) {
        TreeNode* node = nodeQueue.front();
        nodeQueue.pop();
        cin >> item;
        count++;
        if (strcmp(item, "null") != 0) {
            int leftNumber = atoi(item);
            node->left = new TreeNode(leftNumber);
            nodeQueue.push(node->left);
        }
        if (count == n)
            break;
        cin >> item;
        count++;
        if (strcmp(item, "null") != 0) {
            int rightNumber = atoi(item);
            node->right = new TreeNode(rightNumber);
            nodeQueue.push(node->right);
        }
    }
    return root;
}

class Solution {
public:
    TreeNode* pruneTree(TreeNode* root) {
        if (!root) return nullptr;

        // Recursively prune left and right subtrees
        root->left = pruneTree(root->left);
        root->right = pruneTree(root->right);

        // If both left and right subtrees are pruned and the current node's value is 0, prune the current node
        if (!root->left && !root->right && root->val == 0) {
            delete root;
            return nullptr;
        }

        return root;
    }
};

void printTree(TreeNode* root) {
    if (!root) return;
    queue<TreeNode*> nodeQueue;
    nodeQueue.push(root);
    bool hasNonEmptyNode = true;
    while (hasNonEmptyNode) {
        hasNonEmptyNode = false;
        int size = nodeQueue.size();
        for (int i = 0; i < size; ++i) {
            TreeNode* node = nodeQueue.front();
            nodeQueue.pop();
            if (node) {
                cout << node->val << " ";
                if (node->left) {
                    hasNonEmptyNode = true;
                    nodeQueue.push(node->left);
                } else {
                    nodeQueue.push(nullptr);
                }
                if (node->right) {
                    hasNonEmptyNode = true;
                    nodeQueue.push(node->right);
                } else {
                    nodeQueue.push(nullptr);
                }
            } else {
                cout << "null ";
            }
        }
    }
}

int main() {
    TreeNode* root = inputTree();
    Solution sol;
    TreeNode* prunedTree = sol.pruneTree(root);
    printTree(prunedTree);
    return 0;
}