牛客算法周周练D题

//这是一题无向图的题目。
#include<bits/stdc++.h>
 
using namespace std;
 
const int N = 750005;
int n,q;
int h[N],ne[2*N],idx,e[2*N];//无向图记得开二倍int dp[N][3];
void add(int u,int v)
{
    e[idx] = v;
    ne[idx] = h[u];
    h[u] = idx++;
} // 无向图的存储
void dfs(int u, int fa)
{
    f[u] = fa;
    for(int i = h[u];i != -1;i = ne[i])
    {
        int v = e[i];
        if(v == fa) continue;
        dfs(v,u);
    }
}  //遍历所有节点的父亲节点 
int main()
{
    memset(h,-1,sizeof h);
    scanf("%d%d",&n,&q);
    for(int i = 1;i < n;i ++)
    {
        int a,b;
        scanf("%d%d",&a,&b);
        add(a,b),add(b,a);//无向图的存储
    }
    dfs(1,0);
    while(q--)
    {
        int x;
        scanf("%d",&x);
        dp[x][1]++;dp[x][2]++; // 跟自身距离为2的点
        dp[f[x]][0]++;dp[f[f[x]]][0]++; // 父亲节点和爷爷节点
 
        dp[f[x]][1]++;// 与父亲节点距离为1的点
 　　　　 //注意这里不能写dp[x][0]++，因为这个会与父亲节点距离为1的点重复。
        printf("%d
",dp[x][0]+dp[f[x]][1]+dp[f[f[x]]][2]); //输出当前这个点距离为2的点的所有被炸情况
    }
    return 0;
}

题目来源：https://ac.nowcoder.com/acm/contest/5203/D