(leetcode) 2458. 移除子树后的二叉树高度

​​link​​

题意

n个点的树, m次询问. 询问的形式为

割掉以x点为根的子树, 求此时树的最大高度.

首先树的高度一定在叶子节点取到

而以x为根的子树包含的叶子节点设为v1,v2,v3…vx

显然这些叶子节点的dfs序是一个连续的区间.

于是可以直接把所有根节点抽出来, 按dfs序从小到大排成一个数组

对这个数组求前缀max数组pre和后缀max数组suf. 那么由于抽掉的子树中, 叶子节点dfs序区间为[l,r]

答案就是max( pre[l-1], suf[r+1] ), 最后和deep[x]-1求个max. 因为割掉x子树后, 可能x的父节点可以成为新的叶子节点

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    static const int maxn = 3e5+10;
    int l[maxn],r[maxn],dep[maxn],pre[maxn],suf[maxn];
    void dfs(TreeNode* rt,int& id,int depth){
        int u = rt->val, leaf = 0;
        l[u] = 1e9, r[u] = -1e9, dep[u] = depth;
        if( rt->left != nullptr ){
            leaf = 1;
            dfs( rt->left,id,depth+1 );
            int x = rt->left->val;
            l[u] = min( l[u],l[x] );
            r[u] = max( r[u],r[x] );
        }
        if( rt->right != nullptr ){
            leaf = 1;
            dfs( rt->right,id,depth+1 );
            int x = rt->right->val;
            l[u] = min( l[u],l[x] );
            r[u] = max( r[u],r[x] );
        }
        if( !leaf ) l[u] = r[u] = ++id, pre[id] = suf[id] = depth;
    }
    vector<int> treeQueries(TreeNode* root, vector<int>& queries) {
        int id = 0;
        dfs( root,id,0 );
        for(int i=1;i<=id;i++)  pre[i] = max( pre[i],pre[i-1] );
        for(int i=id;i>=1;i--)  suf[i] = max( suf[i],suf[i+1] );
        vector<int>vec;
        for(auto v:queries ){
            int L = l[v], R = r[v];
            int ans = max( dep[v]-1,max( pre[L-1],suf[R+1] ) );
            vec.push_back( ans );
        }
        return vec;
    }
};