leetcode

0 class Solution {

1 public:

2 vector<int> treeQueries(const TreeNode *root, vector<int> &queries) const {

3 static int height[100002];

4 memset(height, 0x00, sizeof(height));

5

6 stack<const TreeNode *> st;

7 st.push(root);

8 while (!st.empty()) {

9 if (st.top() != nullptr) {

10 const auto *root = st.top();

11 st.push(nullptr);

12 if (root->left) st.push(root->left);

13 if (root->right) st.push(root->right);

14 continue;

15 }

16

17 st.pop();

18 const auto *root = st.top();

19 st.pop();

20

21 int &h = height[root->val];

22 if (root->left) h = max(h, height[root->left->val] + 1);

23 if (root->right) h = max(h, height[root->right->val] + 1);

24 }

25

26 queue<const TreeNode *> q;

27 q.push(root);

28 for (int lvl = 0; !q.empty(); lvl++) {

29 int f = 100001, s = 100001;

30

31 queue<const TreeNode *> copy = q;

32 for (int k = size(q); k > 0; k--) {

33 const auto *root = q.front();

34 q.pop();

35 const int val = root->val;

36

37 if (height[val] > height[f])

38 s = f, f = val;

39 else if (height[val] > height[s])

40 s = val;

41

42 if (root->left) q.push(root->left);

43 if (root->right) q.push(root->right);

44 }

45

46 if (size(copy) == 1) {

47 height[copy.front()->val] = lvl - 1;

48 continue;

49 }

50

51 const int hf = height[f], hs = height[s];

52 while (!copy.empty()) {

53 const int n = copy.front()->val;

54 copy.pop();

55 if (n != f)

56 height[n] = hf + lvl;

57 else

58 height[n] = hs + lvl;

59 }

60 }

61

62 for (auto &n : queries)

63 n = height[n];

64 return queries;

65 }

66 };