leetcode

0 class Solution {

1 public:

2 TreeNode *sufficientSubset(TreeNode *root, int limit) const {

3 stack<pair<TreeNode *, int>> st;

4 unordered_set<TreeNode *> deleted;

5

6 st.emplace(root, root->val);

7 while (!st.empty()) {

8 if (st.top().first) {

9 auto [root, sum] = st.top();

10 st.emplace(nullptr, 0);

11 if (root->left) st.emplace(root->left, sum + root->left->val);

12 if (root->right) st.emplace(root->right, sum + root->right->val);

13 continue;

14 }

15

16 st.pop();

17 auto [root, sum] = st.top();

18 st.pop();

19

20 if (!root->left && !root->right) {

21 if (sum < limit) deleted.insert(root);

22 } else {

23 int children = 0, del = 0;

24 if (root->left) {

25 children++;

26 if (deleted.count(root->left)) {

27 root->left = nullptr;

28 del++;

29 }

30 }

31 if (root->right) {

32 children++;

33 if (deleted.count(root->right)) {

34 root->right = nullptr;

35 del++;

36 }

37 }

38 if (children == del) deleted.insert(root);

39 }

40 }

41

42 return !deleted.count(root) ? root : nullptr;

43 }

44 };