leetcode

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 };