bzoj 3572: [Hnoi2014]世界树

  1 #include <bits/stdc++.h>
  2 #define LL long long
  3 #define N 300005
  4 using namespace std;
  5 inline int ra()
  6 {
  7     int x=0,f=1; char ch=getchar();
  8     while (ch<‘0‘ || ch>‘9‘) {if (ch==‘-‘) f=-1; ch=getchar();}
  9     while (ch>=‘0‘ && ch<=‘9‘) {x=x*10+ch-‘0‘; ch=getchar();}
 10     return x*f;
 11 }
 12 int bin[20];
 13 int n,m,q,cnt,tot,K,ind;
 14 int fa[N][20],deep[N],head[N],dfn[N],size[N];
 15 int a[N],b[N],f[N];
 16 int top,st[N],c[N];
 17 int rem[N],belong[N];
 18 struct edge{
 19     int to,next;
 20 }e[N<<1];
 21 void pre_bin(){bin[0]=1; for (int i=1; i<20; i++) bin[i]=bin[i-1]<<1;}
 22 void insert(int x, int y){
 23     if (x==y) return;
 24     e[++cnt].next=head[x]; e[cnt].to=y; head[x]=cnt;
 25     }
 26 void dfs(int x)
 27 {
 28     for (int i=1; bin[i]<=deep[x]; i++)
 29         fa[x][i]=fa[fa[x][i-1]][i-1];
 30     size[x]=1; dfn[x]=++ind;
 31     for (int i=head[x];i;i=e[i].next)
 32     {
 33         if (e[i].to==fa[x][0]) continue;
 34         fa[e[i].to][0]=x;
 35         deep[e[i].to]=deep[x]+1;
 36         dfs(e[i].to);
 37         size[x]+=size[e[i].to];
 38     }
 39 }
 40 int lca(int x, int y)
 41 {
 42     if (deep[x]<deep[y]) swap(x,y);
 43     int t=deep[x]-deep[y];
 44     for (int i=0; bin[i]<=t; i++)
 45         if (bin[i]&t) x=fa[x][i];
 46     for (int i=19; i>=0; i--)
 47         if (fa[x][i]!=fa[y][i])
 48             x=fa[x][i],y=fa[y][i];
 49     return x==y?x:fa[x][0];
 50 }
 51 bool cmp(int a, int b){return dfn[a]<dfn[b];}
 52 void build_vtree()
 53 {
 54     sort(a+1,a+K+1,cmp);
 55     cnt=top=tot=0; int flag=0;
 56     if (belong[1]!=1) st[++top]=1; 
 57     for (int i=1+flag; i<=K; i++)
 58     {
 59         int x=a[i],f=lca(x,st[top]);
 60         if (f==st[top]) {st[++top]=x; continue;}
 61         while (f==lca(st[top-1],x))
 62         {
 63             insert(st[top-1],st[top]); top--;
 64             f=lca(x,st[top]);
 65         }
 66         insert(f,st[top]); st[top]=f; st[++top]=x;
 67     }
 68     while(top>1)insert(st[top-1],st[top]),top--;
 69 }
 70 int dis(int x, int y)
 71 {
 72     return deep[x]+deep[y]-2*deep[lca(x,y)];
 73 }
 74 void dfs1(int x)
 75 {
 76     c[++tot]=x; rem[x]=size[x];
 77     for (int i=head[x];i;i=e[i].next)
 78     {
 79         dfs1(e[i].to);
 80         if (!belong[e[i].to]) continue;
 81         int t1=dis(belong[e[i].to],x),t2=dis(belong[x],x);
 82         if ((t1==t2 && belong[e[i].to]<belong[x]) || t1<t2 || !belong[x])
 83             belong[x]=belong[e[i].to];
 84     }
 85 }
 86 void dfs2(int x)
 87 {
 88     for (int i=head[x];i;i=e[i].next)
 89     {
 90         int t1=dis(belong[x],e[i].to),t2=dis(belong[e[i].to],e[i].to);
 91         if ((t1==t2 && belong[e[i].to]>belong[x]) || t1<t2|| !belong[e[i].to])
 92             belong[e[i].to]=belong[x];
 93         dfs2(e[i].to);
 94     }
 95 }
 96 void query(int a, int b)
 97 {
 98     int x=b,mid=b;
 99     for (int i=18; i>=0; i--)
100         if (deep[fa[x][i]]>deep[a]) x=fa[x][i];
101     rem[a]-=size[x];
102     if (belong[a]==belong[b])
103     {
104         f[belong[a]]+=size[x]-size[b];
105         return;
106     }
107     for (int i=18; i>=0; i--)
108     {
109         int nxt=fa[mid][i];
110         if (deep[nxt]<=deep[a]) continue;
111         int t1=dis(belong[a],nxt),t2=dis(belong[b],nxt);
112         if (t1>t2 || (t2==t1 && belong[b]<belong[a])) mid=nxt;
113     }
114     f[belong[a]]+=size[x]-size[mid];
115     f[belong[b]]+=size[mid]-size[b];
116 }
117 void solve()
118 {
119     K=ra();
120     for (int i=1; i<=K; i++) a[i]=b[i]=ra();
121     for (int i=1; i<=K; i++) belong[a[i]]=a[i];
122     build_vtree(); 
123     dfs1(1); dfs2(1);
124     for (int i=1; i<=tot; i++) 
125         for (int j=head[c[i]];j;j=e[j].next)
126             query(c[i],e[j].to);
127     for (int i=1; i<=tot; i++) 
128         f[belong[c[i]]]+=rem[c[i]];
129     for (int i=1; i<=K; i++) printf("%d ",f[b[i]]);
130     for (int i=1; i<=tot; i++)
131             f[c[i]]=belong[c[i]]=head[c[i]]=rem[c[i]]=0;
132     puts("");//while (1);
133 }
134 int main(int argc, char const *argv[])
135 {
136     n=ra(); pre_bin();
137     for (int i=1; i<n; i++)
138     {
139         int x=ra(),y=ra();
140         insert(x,y); insert(y,x);
141     }
142     dfs(1);
143     memset(head,0,sizeof(head));
144     m=ra();
145     for (int i=1; i<=m; i++)  solve();
146     return 0;
147 }