中规中矩的虚树写法

不难发现，只有同一深度的松鼠才会打架。

那就easy了啊，把同一深度的所有点拎出来，建一棵虚树，合并一下，就好了。

对于一个点，记表示到达该点的松鼠数量，那么

其中表示在原树上的深度。这个方程应该很好理解吧（雾。

然后每次把答案加上就好了(注意要满足先)。

复杂度：瓶颈在求lca上

代码：

#include<bits/stdc++.h>
#define int long long
#define MAXN 200010
using namespace std;
int n,s,A[MAXN],head[MAXN],tot,deep[MAXN],pre[MAXN][20],lg[MAXN],dfsn[MAXN],NODE,stk[MAXN],top,dp[MAXN],ans;
vector<int> P[MAXN];
vector<int> E[MAXN];
struct node{
    int ed,last;
}G[MAXN<<2];
bool cmp(int X,int Y){
    return dfsn[X]<dfsn[Y];
}
void Add(int st,int ed){
    tot++;
    G[tot]=node{ed,head[st]};
    head[st]=tot;
}
int LCA(int x,int y){//求LCA
    if(deep[x]<deep[y])swap(x,y);
    while(deep[x]>deep[y])x=pre[x][lg[deep[x]-deep[y]]-1];
    if(x==y)return x;
    for(int i=lg[deep[x]]-1;i>=0;i--){
        if(pre[x][i]==pre[y][i])continue;
        x=pre[x][i],y=pre[y][i];
    }
    return pre[x][0];
}
void Insert(int x){//虚树中插入节点
    if(top==1){
        if(x!=stk[top])stk[++top]=x;
        return;
    }
    int lca=LCA(stk[top],x);
    while(top>=2&&dfsn[lca]<=dfsn[stk[top-1]]){
        E[stk[top-1]].push_back(stk[top]);
        top--;
    }
    if(lca!=stk[top])E[lca].push_back(stk[top]),stk[top]=lca;
    stk[++top]=x;
}
void DFS(int x,int fa){
    dfsn[x]=++NODE;
    deep[x]=deep[fa]+1;
    pre[x][0]=fa;
    for(int i=1;(1<<i)<=deep[x];i++)pre[x][i]=pre[pre[x][i-1]][i-1];
    for(int i=head[x];i;i=G[i].last){
        int t=G[i].ed;
        if(t==fa)continue;
        DFS(t,x);
    }
}
void solve(int x){//在虚树中dp一波
    const int sz=E[x].size();
    dp[x]=0;
    stk[++top]=x;
    if(sz==0){
        dp[x]=A[x];
        return;
    }
    for(int i=0;i<sz;i++){
        int t=E[x][i];
        solve(t);
        if(dp[t]!=0)dp[x]+=max(1LL,dp[t]-(deep[t]-deep[x]));
    }
}
signed main(){
    for(int i=1;i<=MAXN-10;i++)lg[i]=lg[i-1]+(1<<(lg[i-1])==i);
    scanf("%lld %lld",&n,&s);
    for(int i=1;i<=n;i++)scanf("%lld",&A[i]);
    for(int x,y,i=1;i<=n-1;i++){
        scanf("%lld %lld",&x,&y);
        Add(x,y);
        Add(y,x);
    }
    DFS(s,0);
    for(int i=1;i<=n;i++)P[deep[i]].push_back(i);
    for(int i=1;i<=n;i++){
        sort(P[i].begin(),P[i].end(),cmp);
        if(P[i].size()==0)continue;
        top=0;
        stk[++top]=s;
        const int sz=P[i].size();
        for(int j=0;j<sz;j++){
            int t=P[i][j];
            Insert(t);
        }
        while(top>=2)E[stk[top-1]].push_back(stk[top]),top--;
        top=0;
        solve(s);
        if(dp[s])ans+=max(1LL,dp[s]-1);
        for(int i=1;i<=top;i++)E[stk[i]].clear(),dp[stk[i]]=0;//注意不要用memset，要回滚清空
    }
    cout<<ans;
    return 0;
}