algorithm/Algorithm/cpp/0653.TwoSumIV-InputisaBST/findTarget.cpp at master · wuvei/algorithm · 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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class InorderTreeIterator {
public:
    stack<TreeNode*> st;
    TreeNode* curr;
    bool isReversed;
    InorderTreeIterator(TreeNode* root, bool isReversed) {
        st.push(root);
        this->isReversed = isReversed;
        this->curr = root;
    }

    TreeNode* next() {
        while (curr || !st.empty()) {
            while (curr) {
                st.push(curr);
                curr = isReversed ? curr->right : curr->left;
            }
            curr = st.top();
            st.pop();
            TreeNode* currSave = curr;
            auto right = isReversed ? curr->left : curr->right;
            curr = right;
            return currSave;
        }
        return nullptr;
    }
};

class Solution {
public:
    bool findTarget(TreeNode* root, int k) {
        InorderTreeIterator left(root, false);
        InorderTreeIterator right(root, true);

        TreeNode* l = left.next();
        TreeNode* r = right.next();
        while (l != r) {
            int sum = l->val + r->val;
            if (sum == k) return true;
            else if (sum < k) l = left.next();
            else r = right.next();
        }
        return false;
    }
};