2458.cpp (2004B)
1 class Solution { 2 public: 3 vector<int> treeQueries(const TreeNode *root, vector<int> &queries) const { 4 static int height[100002]; 5 memset(height, 0x00, sizeof(height)); 6 7 stack<const TreeNode *> st; 8 st.push(root); 9 while (!st.empty()) { 10 if (st.top() != nullptr) { 11 const auto *root = st.top(); 12 st.push(nullptr); 13 if (root->left) st.push(root->left); 14 if (root->right) st.push(root->right); 15 continue; 16 } 17 18 st.pop(); 19 const auto *root = st.top(); 20 st.pop(); 21 22 int &h = height[root->val]; 23 if (root->left) h = max(h, height[root->left->val] + 1); 24 if (root->right) h = max(h, height[root->right->val] + 1); 25 } 26 27 queue<const TreeNode *> q; 28 q.push(root); 29 for (int lvl = 0; !q.empty(); lvl++) { 30 int f = 100001, s = 100001; 31 32 queue<const TreeNode *> copy = q; 33 for (int k = size(q); k > 0; k--) { 34 const auto *root = q.front(); 35 q.pop(); 36 const int val = root->val; 37 38 if (height[val] > height[f]) 39 s = f, f = val; 40 else if (height[val] > height[s]) 41 s = val; 42 43 if (root->left) q.push(root->left); 44 if (root->right) q.push(root->right); 45 } 46 47 if (size(copy) == 1) { 48 height[copy.front()->val] = lvl - 1; 49 continue; 50 } 51 52 const int hf = height[f], hs = height[s]; 53 while (!copy.empty()) { 54 const int n = copy.front()->val; 55 copy.pop(); 56 if (n != f) 57 height[n] = hf + lvl; 58 else 59 height[n] = hs + lvl; 60 } 61 } 62 63 for (auto &n : queries) 64 n = height[n]; 65 return queries; 66 } 67 };