牛客-小H和游戏（树的dfs）

原题链接

解题思路：思路见代码，如还有疑惑。请评论区留言

C++代码

#include <bits/stdc++.h>
// #define int long long 
using namespace std;
const int N = 8e5 + 10, M = N * 2;
int h[N], e[M], ne[M], idx;
int f[N], n, q;
int cnt[N][3]; //0表示N点自己、1表示N点儿子、2表示N点孙子

void add(int a, int b)  // 添加一条边a->b
{
    e[idx] = b, ne[idx] = h[a], h[a] = idx ++ ;
}

void dfs(int u, int fa)
{
    f[u] = fa;
    for (int i = h[u]; ~i; i = ne[i])
    {
        int j = e[i];
        if (j == fa) continue;
        dfs(j, u);
    }
}

signed main()
{
    memset(h, -1, sizeof h);
    cin >> n >> q;
    for (int i = 1; i < n; i ++ )
    {
        int a, b;
        cin >> a >> b;
        add(a, b), add(b, a);
    }
    dfs(1, 0);
    while (q -- )
    {
        int x;
        cin >> x;
        ++ cnt[f[f[x]]][0]; //波及到爷爷
        ++ cnt[f[x]][0]; //波及到父亲
        ++ cnt[f[x]][1]; //波及到兄弟及自己
        ++ cnt[x][1]; //波及到儿子
        ++ cnt[x][2]; //波及到孙子
        
        //cnt[x][0]表示自己波及到自己、儿子波及到自己和孙子波及到自己
        //cnt[f[f[x]]]表示爷爷波及到自己
        //cnt[f[x]][1]表示父亲波及到自己和兄弟波及到自己
        cout << cnt[x][0] + cnt[f[f[x]]][2] + cnt[f[x]][1] << endl;
    }
    return 0;
}