bzoj3631: [JLOI2014]松鼠的新家（LCA+差分）

 1 #include<stdio.h>
 2 #include<algorithm>
 3 using namespace std;
 4 const int maxn = 300010;
 5 struct node{
 6     int to,next;
 7 }e[maxn*2];
 8 int n,a[maxn],head[maxn],dep[maxn],fa[maxn][20],w[maxn],f[maxn],tot,logn;
 9 
10 void insert(int u, int v){
11     e[++tot].to=v; e[tot].next=head[u]; head[u]=tot;
12 }
13 
14 void dfs(int u, int f, int d){
15     dep[u]=d; fa[u][0]=f;
16     for (int i=1; i<=logn; i++) fa[u][i]=fa[fa[u][i-1]][i-1];
17     for (int i=head[u]; i; i=e[i].next)
18         if (e[i].to!=f) dfs(e[i].to,u,d+1);
19 }
20 
21 void lca(int u, int v){
22     if (dep[u]<dep[v]) swap(u,v);
23     int x=u,y=v;
24     while (dep[u]>dep[v]){
25         for (int i=logn; i>=0; i--)
26             if (dep[fa[u][i]]>dep[v])
27                 u=fa[u][i];
28         u=fa[u][0];
29     }
30     if (u==v){  //在某条链上 
31         w[fa[u][0]]--;
32         w[x]++;
33         return;
34     }
35     for (int i=logn; i>=0; i--)  //在分叉口上 
36         if (fa[u][i]!=fa[v][i]){
37             u=fa[u][i];
38             v=fa[v][i];
39         }
40     u=fa[u][0];
41     w[x]++; w[y]++;
42     w[u]--; w[fa[u][0]]--;
43     return; 
44 }
45 
46 void get_ans(int u){
47     f[u]=w[u];// printf("  %d\n", w[u]);
48     for (int i=head[u],v; i; i=e[i].next){
49         if ((v=e[i].to)==fa[u][0]) continue;
50         get_ans(e[i].to);
51         f[u]+=f[v];
52     }
53 }
54 
55 int main(){
56     scanf("%d", &n);
57     for (int i=1; i<=n; i++) scanf("%d", &a[i]); tot=1; while ((1<<logn)<n) logn++;
58     for (int i=1,u,v; i<n; i++) scanf("%d%d", &u, &v),insert(u,v),insert(v,u); insert(0,a[1]);
59     dfs(0,0,1);
60     for (int i=2; i<=n; i++){
61         lca(a[i-1],a[i]);
62     }
63     get_ans(a[1]);
64     for (int i=1; i<=n; i++) printf("%d\n", (i==a[1])?f[i]:f[i]-1);
65     return 0;
66 }