洛谷 P3258 [JLOI2014]松鼠的新家（树上差分）

题目链接：https://www.luogu.com.cn/problem/P3258

对点差分：

sa[u]++;
sa[v]++;
sa[lca(u,v)]--;
sa[f[lca(u,v)][0]]--;

对边差分：

sa[u]++;
sa[v]++;
sa[lca(u,v)]--;

注意这道题a[2]~a[n]，它们既被作为前一段的终点，同时也是下一段的起点，所以每个点差分后会被多算一次，需要减掉。

AC代码：

#include<cstdio>
#include<iostream>
#include<cstring>
#include<algorithm>
using namespace std;
const int N=3*1e5+5;
int n,tot;
int head[N],dep[N],f[N][30],a[N],sa[N];
struct node{
    int to,next;
}edge[N<<1];
void init(){
    memset(head,-1,sizeof(head));
}
void add(int u,int v){
    edge[tot].to=v;
    edge[tot].next=head[u];
    head[u]=tot++;
}
void DFS(int u,int fa){
    f[u][0]=fa;
    dep[u]=dep[fa]+1;
    for(int i=1;(1<<i)<=dep[u];i++) f[u][i]=f[f[u][i-1]][i-1];
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].to;
        if(v==fa) continue;
        DFS(v,u);
    }
}
int lca(int u,int v){
    if(dep[u]<dep[v]) swap(u,v);
    for(int i=20;i>=0;i--) if(dep[f[u][i]]>=dep[v]) u=f[u][i];
    if(u==v) return u;
    for(int i=20;i>=0;i--) if(f[u][i]!=f[v][i]) u=f[u][i],v=f[v][i];
    return f[u][0];
}
void get_sa(int u,int fa){
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].to;
        if(v==fa) continue;
        get_sa(v,u);
        sa[u]+=sa[v];
    }
}
int main(){
    init();
    scanf("%d",&n);
    for(int i=1;i<=n;i++) scanf("%d",&a[i]);
    for(int i=1;i<n;i++){
        int x,y;
        scanf("%d%d",&x,&y);
        add(x,y); add(y,x);
    }
    DFS(1,0);
    for(int i=2;i<=n;i++){
        int u,v;
        u=a[i-1],v=a[i];
        sa[u]++;
        sa[v]++;
        int x=lca(u,v);
        sa[x]--;
        sa[f[x][0]]--;
    }
    get_sa(1,0);
    for(int i=2;i<=n;i++) sa[a[i]]--;
    for(int i=1;i<=n;i++) printf("%d
",sa[i]);
    return 0;
}