考虑到一个节点的儿子孙子和兄弟可能有多个,但是父节点和祖父节点只有一个,所以我们把不断累加的答案放进它的父节点和祖父节点中去,然后在需要计算时,再通过父节点和祖父节点,把答案加回来即可。
定义sum[u][i]:u节点受到波及的次数,且可以向外延伸i个单位长度。
对于每次轰炸,我们进行如下操作:
1.sum[u][2]++,sum[f[u]][1]++,sum[f[f[u]][0]++。 (若无父节点或祖父节点则省略)
2.ans=sum[u][2]+sum[u][1]+sum[u][0] (孙子,儿子,和自己给u节点带来的影响)
3.ans+=sum[f[f[u]]][2] (祖父节点带来的影响,如无祖父节点则省略)
4.ans+=sum[f[u]][2] (父节点带来的影响,若无父节点则省略)
5.ans+=sum[f[u]][1]-sum[u][2] (兄弟节点带来的影响,若无兄弟节点则省略。这里为什么需要减去sum[u][2]呢,因为u之前的操作sum[f[u]][1]++,导sum[f[u]][1]多加了sum[u][2]次,所以需要减去,即:sum[f[u]][1]中包含了sum[u][2]它自己的贡献,而它自己,不是它的兄弟!)
#include <bits/stdc++.h>
#define int long long
using namespace std;
const int N=1e6+5;
int n,q,u,v,ans;
int sum[N][3],f[N];
int cnt,head[N];
struct edge{int next,to;}e[N<<1];
inline void add(int u,int v)
{
cnt++;
e[cnt].next=head[u];
e[cnt].to=v;
head[u]=cnt;
}
void dfs(int u,int fa)
{
for (register int i=head[u]; i; i=e[i].next)
if (e[i].to!=fa)
{
f[e[i].to]=u;
dfs(e[i].to,u);
}
}
signed main(){
scanf("%lld%lld",&n,&q);
for (register int i=1; i<n; ++i) scanf("%lld%lld",&u,&v),add(u,v),add(v,u);
dfs(1,0);
while (q--)
{
scanf("%lld",&u);
sum[u][2]++; if (f[u]) sum[f[u]][1]++; if (f[f[u]]) sum[f[f[u]]][0]++;
ans=0;
ans=sum[u][2]+sum[u][1]+sum[u][0];
if (f[u]) ans+=sum[f[u]][2],ans+=sum[f[u]][1]-sum[u][2];
if (f[f[u]]) ans+=sum[f[f[u]]][2];
printf("%lld\n",ans);
}
return 0;
}