BZOJ3631: [JLOI2014]松鼠的新家 树链剖分/lca/树上查分

#include<bits/stdc++.h>
#define rep(i,l,r) for(int i=l;i<=r;i++)
#define N 310000
using namespace std;
int head[N],dep[N],sz[N],son[N],fa[N],a[N],top[N],tot,n,p,q,s[N],ans[N],start[N],end[N];
struct node {
     int to,next;
}e[N<<1];
inline void ins(int u,int v) {
     e[++tot].to=v; e[tot].next=head[u]; head[u]=tot;
}
inline int in() {
     int x=0,ch=getchar();
     while(ch<‘0‘ || ch>‘9‘) ch=getchar();
     while(ch>=‘0‘ && ch<=‘9‘) x=x*10+ch-48,ch=getchar();
     return x;
}
void dfs(int x,int pre) {
     dep[x]=dep[fa[x]=pre]+1; sz[x]=1;
     for(int k=head[x];k;k=e[k].next) if(e[k].to!=pre){
          dfs(e[k].to,x);
          sz[x]+=sz[e[k].to];
          if(sz[e[k].to]>sz[son[x]]) son[x]=e[k].to;
     }
}
void dfs1(int x,int chain) {
     top[x]=chain; 
     if(son[x]) dfs1(son[x],chain);
     for(int k=head[x];k;k=e[k].next) if(e[k].to!=fa[x]&&e[k].to!=son[x]) dfs1(e[k].to,e[k].to);
}
int lca(int x,int y) {
     int a=top[x],b=top[y];
     while(a!=b) {
          if(dep[a]<dep[b]) swap(a,b),swap(x,y);
          x=fa[a]; a=top[x];
     }
     if(dep[x]<dep[y])return x; else return y;
}
int getans(int x) {
     int sum=start[x];
     for(int k=head[x];k;k=e[k].next) if(e[k].to!=fa[x]) sum+=getans(e[k].to);
     sum-=end[x];
     return ans[x]=sum;
}
int main () {
    n=in();
    rep(i,1,n) a[i]=in();
    rep(i,1,n-1) p=in(),q=in(),ins(p,q),ins(q,p); ins(0,1); ins(1,0);
    dfs(1,0);  dfs1(1,1); 
    rep(i,1,n-1) {
        int c=lca(a[i],a[i+1]);
        start[a[i]]++; start[a[i+1]]++; end[c]++; end[fa[c]]++;
    }
    getans(1);
    rep(i,1,n) printf("%d\n",ans[i]-(i!=a[1]));
}