leetcode/Problems/0437.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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
// Brute force
class Solution {
  public:
    int pathSum(TreeNode *root, int targetSum) {
        if (!root) return 0;

        stack<TreeNode *> st;
        queue<TreeNode *> q;

        st.push(root);
        while (!st.empty()) {
            TreeNode *root = st.top();
            st.pop();
            q.push(root);
            if (root->left) st.push(root->left);
            if (root->right) st.push(root->right);
        }

        int res = 0;
        while (!q.empty()) {
            stack<pair<TreeNode *, long long>> st;
            st.push({q.front(), 0});
            q.pop();
            while (!st.empty()) {
                auto [root, sum] = st.top();
                st.pop();
                sum += root->val;
                if (sum == targetSum) res++;
                if (root->left) st.push({root->left, sum});
                if (root->right) st.push({root->right, sum});
            }
        }

        return res;
    }
};

// Optimized
class Solution {
  public:
    int pathSum(TreeNode *root, int targetSum) {
        if (!root) return 0;
        queue<pair<TreeNode *, vector<long long>>> q;
        int res = 0;

        q.push({root, {}});
        while (!q.empty()) {
            auto &[root, vec] = q.front();
            long long sum = root->val + (vec.size() ? vec.back() : 0);

            for (int num : vec)
                if (sum - num == targetSum) res++;
            if (sum == targetSum) res++;

            if (root->left) {
                q.push({root->left, vec});
                q.back().second.push_back(sum);
            }

            if (root->right) {
                q.push({root->right, vec});
                q.back().second.push_back(sum);
            }

            q.pop();
        }

        return res;
    }
};