小H和游戏

链接：

来源：牛客网

https://ac.nowcoder.com/acm/contest/5203/D

题目描述

小H正在玩一个战略类游戏，她可以操纵己方的飞机对敌国的N座城市(编号为1~N)进行轰炸
敌国的城市形成了一棵树，小H会依次进行Q次轰炸，每次会选择一个城市A进行轰炸，和这座城市距离不超过2的城市都会受损(这里距离的定义是两点最短路径上的边数)，轰炸结束后，小H还想知道当前城市A受损的次数
作为游戏的开发者之一，你有义务回答小H的问题

输入描述:

第1行，两个整数N(1≤N≤750000)、Q(1≤Q≤750000)
第2~N行，每行两个整数表示树上的一条边第N+1~N+Q行，每行一个整数，表示小H这次轰炸的城市

输出描述:

输出Q行，每行一个整数表示这一次轰炸的城市在此次轰炸后共计受损几次

示例1

输入

[复制]

4 4 
1 2 
2 3 
3 4 
1 
2 
3 
4

输出

[复制]

1 
2 
3 
3

思路

每个点存三个信息，分别是自己被直接炸的次数，儿子节点被直接炸的次数，孙子节点被直接炸的次数，每次自己被炸时，只取更新父亲和爷爷节点的信息。查询时，一个点被炸得次数有自己,儿子,孙子,父亲,兄弟,爷爷节点被直接炸得次数构成。

除了兄弟节点的信息其他都容易统计，自己和兄弟节点都是父亲节点的儿子，所以每次不计自己被炸得次数，直接加上上父亲节点的儿子节点被炸次数。注意特判根节点，根节点没有父亲，所以还是需要加上自己被炸的次数。

代码

#include<bits/stdc++.h>
using namespace std;
const int N=800000;
int h[N],idx,fa[N],g[N][3],n,q;//0自己被点 1距离为1 2距离为2
struct eg{
    int v,nex;
}e[N*2];
void add(int u,int v){
    e[++idx]=(eg){v,h[u]};
    h[u]=idx;
}
void dfs(int u,int pre){
    fa[u]=pre;
    for(int i=h[u];~i;i=e[i].nex){
        int v=e[i].v;
        if(v==pre) continue;
        dfs(v,u);
    }
}
int solve (int u){
    int res=0,fu=fa[u],ffu=fa[fa[u]];
    if(fu!=0)
        res+=g[fu][1]+g[fu][0];
    if(ffu!=0)
        res+=g[ffu][0];
    if(fu!=0)
        g[fu][1]++;
    if(ffu!=0)
        g[ffu][2]++;
    return res;
}
void solve(){
    int x;
    while(q--){
        scanf("%d",&x);
        g[x][0]++;
        int res=g[x][2]+g[x][1]+1;
        res+=solve(x);
        if(fa[x]==0) res+=g[x][0]-1;
        printf("%d
",res);
    }
}
int main(){
    int u,v;
    scanf("%d%d",&n,&q);
    memset(h,-1,sizeof h);
    idx=0;
    n--;
    while(n--){
        scanf("%d%d",&u,&v);
        add(u,v);
        add(v,u);
    }
    dfs(1,0);
    solve();
    return 0;
}