luoguP3258 [JLOI2014]松鼠的新家 题解(树上差分)

P3258 [JLOI2014]松鼠的新家  题目

树上差分:树上差分总结

#include<iostream>
#include<cstdlib>
#include<cstdio>
#include<cmath>
#include<cstring>
#include<algorithm>
#include<iomanip>
#include<ctime>
#include<queue>
#include<stack>
#define rg register
#define lst long long
#define N 300050
using namespace std;

int n,K,cnt,ans;
struct EDGE{
    int to,nxt;
}edge[N<<];
int head[N],an[N];
int fa[N],v[N],deep[N];
int f[N][];

inline int read()
{
    rg int s=,m=;rg char ch=getchar();
    while(ch!='-'&&(ch<''||ch>''))ch=getchar();
    if(ch=='-')m=-,ch=getchar();
    while(ch>=''&&ch<='')s=(s<<)+(s<<)+ch-'',ch=getchar();
    return s*m;
}

inline void add(rg int p,rg int q)
{
    edge[++cnt]=(EDGE){q,head[p]};
    head[p]=cnt;
}

void dfs(rg int now,rg int fm,rg int dep)
{
    fa[now]=fm;deep[now]=dep;
    f[now][]=fa[now];
    for(rg int i=;i<=;++i)
        f[now][i]=f[f[now][i-]][i-];
    for(rg int i=head[now];i;i=edge[i].nxt)
    {
        rg int qw=edge[i].to;
        if(qw!=fm)
            dfs(qw,now,dep+);
    }
}

inline int lca(rg int p,rg int q)
{
    if(deep[p]<deep[q])swap(p,q);
    while(deep[p]>deep[q])
        for(rg int i=;i>=;--i)
            if(deep[f[p][i]]>=deep[q])
                p=f[p][i];
    while(p!=q)
    {
        for(rg int i=;i>=;--i)
            if(f[p][i]!=f[q][i])
                p=f[p][i],q=f[q][i];
        if(fa[p]==fa[q])
            p=q=fa[p];
    }
    return p;
}

inline void Insert(rg int p,rg int q)
{
    v[p]++,v[q]++;
      rg int LCA=lca(p,q);
//    printf("lca(%d,%d)=%d\n",p,q,LCA);
      v[LCA]--,v[fa[LCA]]--;
}

void sum(rg int now)
{
    for(rg int i=head[now];i;i=edge[i].nxt)
    {
        rg int qw=edge[i].to;
        if(qw!=fa[now])
        {
            sum(qw);
            v[now]+=v[qw];
        }
    }
}

int main()
{
    n=read();
    for(rg int i=;i<=n;++i)an[i]=read();
    for(rg int i=;i<n;++i)
    {
        rg int p=read(),q=read();
        add(p,q),add(q,p);
    }
      dfs(,,);
    for(rg int i=;i<n;++i)
    {
        Insert(an[i],an[i+]);
    }
//    for(rg int i=\1;i<=n;++i)cout<<deep[i]<<" ";cout<<endl;
/*    for(rg int i=1;i<=n;++i)
    {
        for(rg int j=0;j<=2;j++)
        {
            printf("f[%d][%d]=%d ",i,j,f[i][j]);
        }
        cout<<endl;
    }
    cout<<endl;*/
    add(,);
    sum();
    for(rg int i=;i<=n;++i)v[an[i]]--;
    for(rg int i=;i<=n;++i)printf("%d\n",v[i]);
    return ;
}