leetcode/Problems/1530.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
// Recursive
class Solution {
  public:
    int countPairs(TreeNode *root, int distance) {
        unordered_map<TreeNode *, vector<int>> um;
        stack<TreeNode *> st;
        int res = 0;

        st.push(root);
        while (!st.empty()) {
            TreeNode *root = st.top();
            if (root) {
                st.push(nullptr);
                if (!root->left && !root->right)
                    um[root].push_back(1);
                else {
                    if (root->left) st.push(root->left);
                    if (root->right) st.push(root->right);
                }
                continue;
            }
            st.pop();
            root = st.top();
            st.pop();

            for (const int n : um[root->right])
                um[root].push_back(n + 1);

            for (const int a : um[root->left]) {
                um[root].push_back(a + 1);
                for (const int b : um[root->right])
                    if (a + b <= distance) res++;
            }
        }
        return res;
    }
};

// Iterative
class Solution {
    int res = 0;
    vector<int> rec(TreeNode *root, int distance) {
        if (!root->left && !root->right) return {1};
        vector<int> left, right, sum;
        if (root->left) left = rec(root->left, distance);
        if (root->right) right = rec(root->right, distance);

        sum.reserve(left.size() + right.size());
        for (const int b : right)
            sum.push_back(b + 1);
        for (const int a : left) {
            sum.push_back(a + 1);
            for (const int b : right) {
                res += (a + b <= distance);
            }
        }
        return sum;
    }

  public:
    int countPairs(TreeNode *root, int distance) {
        rec(root, distance);
        return res;
    }
};