bzoj3631 [JLOI2014]松鼠的新家——树上差分-CSDN博客

题目：https://www.lydsy.com/JudgeOnline/problem.php?id=3631

树上差分；注意路径的结尾被多算了一次，最后要减去(不能提前减)。

代码如下：

#include<iostream>
#include<cstdio>
#include<cstring>
using namespace std;
int const maxn=300005;
int n,a[maxn],head[maxn],ct,f[maxn],re[maxn],fa[maxn][22],bin[22],deep[maxn];
struct N{
    int to,next;
    N(int t=0,int n=0):to(t),next(n) {}
}edge[maxn<<1];
bool vis[maxn];
void dfs(int x)
{
    for(int i=1;i<=18;i++)
    {
        if(deep[x]>=bin[i])fa[x][i]=fa[fa[x][i-1]][i-1];
        else break;
    }    
    for(int i=head[x],u;i;i=edge[i].next)
    {
        if(edge[i].to==fa[x][0])continue;
        deep[u=edge[i].to]=deep[x]+1;
        fa[u][0]=x;
        dfs(u);
    }
}
int lca(int x,int y)
{
    if(deep[x]>deep[y])swap(x,y);
    int t=deep[y]-deep[x];
    for(int i=0;i<=18;i++)
        if(t&bin[i])y=fa[y][i];
    for(int i=18;i>=0;i--)
        if(fa[x][i]!=fa[y][i])
            x=fa[x][i],y=fa[y][i];
    if(x==y)return y;
    return fa[y][0];
}
void dfs2(int x)
{
//    cout<<x<<endl;
    if(vis[x])return;
    vis[x]=1;
    for(int i=head[x],u;i;i=edge[i].next)
    {
        u=edge[i].to;
        if(u==fa[x][0])continue;
        dfs2(u);
        f[x]+=f[u];
    }
}
int main()
{
    bin[0]=1;
    for(int i=1;i<20;i++)bin[i]=(bin[i-1]<<1);
    scanf("%d",&n);
    for(int i=1;i<=n;i++)scanf("%d",&a[i]);
    for(int i=1,x,y;i<n;i++)
    {
        scanf("%d%d",&x,&y);
        edge[++ct]=N(y,head[x]);head[x]=ct;
        edge[++ct]=N(x,head[y]);head[y]=ct;
    }
    dfs(1);
    for(int i=2;i<=n;i++)
    {
        f[a[i-1]]++;f[a[i]]++;re[a[i]]++;
        int l=lca(a[i-1],a[i]);
        f[l]--;f[fa[l][0]]--;
    }
    dfs2(1);
    for(int i=1;i<=n;i++)
        printf("%d\n",f[i]-re[i]);
    return 0;
}