leetcode/Problems/0814.cpp at master · DimitrijeDobrota/leetcode · GitHub

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
29
30
31
32
class Solution {
  public:
    TreeNode *pruneTree(TreeNode *root) {
        TreeNode dummy(1, root, nullptr);
        unordered_set<TreeNode *> has;
        stack<TreeNode *> st;
        st.push(&dummy);
        while (!st.empty()) {
            TreeNode *root = st.top();
            if (!root) {
                st.pop();
                root = st.top(), st.pop();
                if (has.count(root->left))
                    has.insert(root);
                else
                    root->left = nullptr;

                if (has.count(root->right))
                    has.insert(root);
                else
                    root->right = nullptr;

                if (root->val == 1) has.insert(root);
                continue;
            }
            st.push(nullptr);
            if (root->left) st.push(root->left);
            if (root->right) st.push(root->right);
        }
        return dummy.left;
    }
};