diff --git a/2458. Height of Binary Tree After Subtree Removal Queries b/2458. Height of Binary Tree After Subtree Removal Queries new file mode 100644 index 0000000..02239db --- /dev/null +++ b/2458. Height of Binary Tree After Subtree Removal Queries @@ -0,0 +1,48 @@ +class Solution { +public: + vector< pair > h = vector>(100001, {0, 0}); + pair solve(TreeNode* root, int level){ + if(root == nullptr) return {0, 0}; + if(root->left == nullptr && root->right == nullptr){ + return h[root->val] = {1, level}; + } + //traversing the tree and updating the h vector + solve(root->left, level + 1); + solve(root->right, level + 1); + + h[root->val] = {max( + root->right != nullptr ? h[root->right->val].first : 0, + root->left != nullptr ? h[root->left->val].first : 0 + ) + 1, level}; + return h[root->val]; + } + vector treeQueries(TreeNode* root, vector& queries) { + solve(root, 0); + vector ans; + vector v[100000]; + int longestPath = 0; + for(auto x : h){ // filling the v vector while calculating longestPath + if(x.second == 0) continue; + longestPath = max(longestPath, x.first); + v[x.second].push_back(x.first); + } + + // Sort to get the longest && the second longest path of each level + for(auto &x : v) sort(x.begin(), x.end(), greater()); + + for(auto &q : queries){ + int node = q; + int level = h[node].second; + if(h[node].first == v[level][0] && v[level].size() >= 2){ + q = longestPath - v[level][0] + v[level][1]; + continue; + }else if(h[node].first == v[level][0]){ + q = longestPath - v[level][0]; + continue; + } + q = longestPath; + } + + return queries; + } +};