leetcode

Solution to some Leetcode problems written in C++

git clone git://git.dimitrijedobrota.com/leetcode.git

1530.cpp (1797B)

0 // Recursive 1 class Solution { 2 public: 3 int countPairs(TreeNode *root, int distance) { 4 unordered_map<TreeNode *, vector<int>> um; 5 stack<TreeNode *> st; 6 int res = 0; 7 8 st.push(root); 9 while (!st.empty()) { 10 TreeNode *root = st.top(); 11 if (root) { 12 st.push(nullptr); 13 if (!root->left && !root->right) 14 um[root].push_back(1); 15 else { 16 if (root->left) st.push(root->left); 17 if (root->right) st.push(root->right); 18 } 19 continue; 20 } 21 st.pop(); 22 root = st.top(); 23 st.pop(); 24 25 for (const int n : um[root->right]) 26 um[root].push_back(n + 1); 27 28 for (const int a : um[root->left]) { 29 um[root].push_back(a + 1); 30 for (const int b : um[root->right]) 31 if (a + b <= distance) res++; 32 } 33 } 34 return res; 35 } 36 }; 37 38 // Iterative 39 class Solution { 40 int res = 0; 41 vector<int> rec(TreeNode *root, int distance) { 42 if (!root->left && !root->right) return {1}; 43 vector<int> left, right, sum; 44 if (root->left) left = rec(root->left, distance); 45 if (root->right) right = rec(root->right, distance); 46 47 sum.reserve(left.size() + right.size()); 48 for (const int b : right) 49 sum.push_back(b + 1); 50 for (const int a : left) { 51 sum.push_back(a + 1); 52 for (const int b : right) { 53 res += (a + b <= distance); 54 } 55 } 56 return sum; 57 } 58 59 public: 60 int countPairs(TreeNode *root, int distance) { 61 rec(root, distance); 62 return res; 63 } 64 };