bzoj_世界树_虚树_树形DP_lca_倍增

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cmath>
#define INF 0x3f3f3f3f
#define rep0(i, n) for (int i = 0; i < n; i++)
#define rep1(i, n) for (int i = 1; i <= n; i++)
#define rep_0(i, n) for (int i = n - 1; i >= 0; i--)
#define rep_1(i, n) for (int i = n; i > 0; i--)
#define MAX(x, y) (((x) > (y)) ? (x) : (y))
#define MIN(x, y) (((x) < (y)) ? (x) : (y))
#define mem(x, y) memset(x, y, sizeof(x))
#define MAXN 300010
#define MAXL 20

using namespace std;
typedef long long ll;
typedef pair<int, int> pp;
struct Edge
{
    int to, nxt;

} edges[MAXN * 2];
int head[MAXN], tot, n, q, pa[MAXN][MAXL], sz[MAXN], dep[MAXN];
int clc[MAXN], cl, bk[MAXN];
int qu[MAXN], tail, qu0[MAXN];
int st[MAXN], top, root;
pp nbr[MAXN];
int ans[MAXN];

bool cmp(int x, int y)
{
    return clc[x] < clc[y];
}
void dfs(int u, int fa, int d)
{
    dep[u] = d;
    pa[u][0] = fa;
    clc[u] = cl++;
    int cnt = 1;

    for (int i = head[u]; i; i = edges[i].nxt)
    {
        int v = edges[i].to;
        if (v == fa)
            continue;
        dfs(v, u, d + 1);
        cnt += sz[v];

    }
    sz[u] = cnt;   //预处理出子树size

}
void init()
{
    dfs(1, 0, 0);
    for (int k = 0; k + 1 < MAXL; k++)
    {
        for (int i = 1; i <= n; i++)
        {
            if (pa[i][k] > 0)
            {
                pa[i][k + 1] = pa[pa[i][k]][k];     //倍增求lca 以及便于求虚树中节点在原树中的父节点

            }
        }
    }
}
void addEdge(int u, int v)  // 单向加边建虚树 加速
{

    edges[++tot].to = v;    // 邻接表以 i = 0 为边界  边应该从1开始存
    edges[tot].nxt = head[u];
    head[u] = tot;

}
int lca(int a, int b)  
{
    if (dep[a] > dep[b])
        swap(a, b);

    for (int k = 0; k < MAXL; k++)
    {
        if ((dep[b] - dep[a]) >> k & 1)
            b = pa[b][k];

    }
    if (a == b)
        return a;


    for (int k = MAXL - 1; k >= 0; k--)
    {
        if (pa[a][k] == pa[b][k])
            continue;
        a = pa[a][k];
        b = pa[b][k];

    }
    return pa[a][0];

}
void build()
{
    top = 0;     //初始化栈
    tot = 0;     // 初始化后原树的邻接表信息废弃  初始化边数组以建立虚树
    int u, fa;
    for (int i = 0; i < tail; i++)
    {
        u = qu[i];

        if (!top)
        {
            st[top++] = u;

        }
        else
        {
            fa = lca(u, st[top - 1]);     // 与栈顶节点的lca

            while (top > 1 && dep[st[top - 2]] >= dep[fa])
            {
                addEdge(st[top - 2], st[top - 1]);
                top--;


            }
            if (dep[fa] < dep[root])   // 建立的虚树上不一定有1节点    需要求深度最小的节点作为root
            {
                root = fa;       
            }
            if (top && dep[st[top - 1]] > dep[fa])    // 如果lca不在栈中 先清空fa的邻接表信息后加边
            {

                head[fa] = 0;
                nbr[fa].second = INF;   // 新入虚树节点  初始化其与 距离最近的点 的距离
                addEdge(fa, st[top - 1]);
                st[top - 1] = fa;
            }
            else if (!top)
            {

                head[fa] = 0;
                nbr[fa].second = INF;
                st[top++] = fa;


            }
            st[top++] = u;


        }
    }
    for (int i = 0; i < top - 1; i++)
        addEdge(st[i], st[i + 1]);

}

void dp(int u, int fa)   //第一遍树形dp求各节点子树范围内 与距离最近点的 距离
{
    int tmp, dis = INF;


    for (int i = head[u]; i; i = edges[i].nxt)
    {
        int v = edges[i].to;
        if (v == fa)
            continue;
        dp(v, u);

        if (dis > nbr[v].second + dep[v] - dep[u])
        {
            tmp = nbr[v].first;
            dis = nbr[v].second + dep[v] - dep[u];

        }
        else if (dis == nbr[v].second + dep[v] - dep[u] && nbr[v].first < tmp)    // 有两个相同距离点的情况
        {
            tmp = nbr[v].first;
        }


    }
    if (bk[u] != q)
    {
        nbr[u].first = tmp;
        nbr[u].second = dis;
    }
}
void dp1(int u, int fa)    // 第二次dp求父亲节点范围内 的距离
{
    if (fa > 0 && bk[u] != q)
    {
        if (nbr[u].second > nbr[fa].second + dep[u] - dep[fa] || (nbr[u].second == nbr[fa].second + dep[u] - dep[fa] && nbr[u].first > nbr[fa].first))
        {
            nbr[u].first = nbr[fa].first;
            nbr[u].second = nbr[fa].second + dep[u] - dep[fa];
        }

        ans[nbr[u].first]++;    // 求出最近点后直接计数


    }
    else if (bk[u] != q)
        ans[nbr[u].first]++;
    for (int i = head[u]; i; i = edges[i].nxt)
    {
        int v = edges[i].to;
        if (v == fa)
            continue;

        dp1(v, u);
    }

}
int solve(int u, int fa)    // 第三次搜索 求不在虚树内的节点的最近节点
{
    
    /**
    由在虚树内的节点向上找 “原树内的父节点”  求出
    虚树内相连两节点 之间的 “原树内的父节点” 的最近关键点 可有这两个节点的信息判断
    */
    int cnt = sz[u];

    if (fa > 0 && dep[u] - dep[fa] > 1)
    {
        int step, tmp, v = u, v1;
        if (nbr[u].first == nbr[fa].first)
        {
            step = dep[u] - dep[fa] - 1;

        }
        else
        {
            tmp = dep[u] - dep[fa] - nbr[u].second + nbr[fa].second;
            step = tmp / 2;
            if (!(tmp & 1) && nbr[u].first > nbr[fa].first)
                step--;

        }

        for (int k = MAXL - 1; k >= 0; k--)
        {
            if (step & (1 << k))
                v = pa[v][k];

        }



        ans[nbr[u].first] += sz[v] - sz[u];
        //cnt = sz[v];
        v1 = v;
        step = dep[v] - dep[fa] - 1;
        if (step > 0)
        for (int k = MAXL - 1; k >= 0; k--)
        {
            if (step & (1 << k))
                v1 = pa[v1][k];

        }
        cnt = sz[v1];

        if (nbr[u].first != nbr[fa].first)
            ans[nbr[fa].first] += sz[v1] - sz[v];



    }
    else if (fa < 0 && u != 1)   // root可能不是1  此时要单独计数root父节点范围内的节点数
    {
        ans[nbr[u].first] += sz[1] - sz[u];

    }

    int tmp = 1;
    for (int i = head[u]; i; i = edges[i].nxt)
    {
        int v = edges[i].to;
        if (v == fa)
            continue;

        tmp += solve(v, u);


    }

    ans[nbr[u].first] += sz[u] - tmp;     // 树上的点上可能有 不在虚树内的子节点 


    return cnt;     // 虚树节点 -> 虚树子节点  一个子树方向  原树内的真实节点数





}

int main()
{
    #ifndef ONLINE_JUDGE
        freopen("in.txt", "r", stdin);
    #endif // ONLINE_JUDGE
    scanf("%d", &n);
    int u, v;

    for (int i = 0; i < n - 1; i++)
    {
        scanf("%d %d", &u, &v);
        addEdge(u, v);
        addEdge(v, u);

    }



    init();

    scanf("%d", &q);

    while (q)
    {
        int t, tmp;
        scanf("%d", &t);
        tail = t;
        for (int i = 0; i < t; i++)
        {
            scanf("%d", &tmp);

            bk[tmp] = q;
            qu[i] = tmp;
            qu0[i] = tmp;

            nbr[tmp].first = tmp;
            nbr[tmp].second = 0;
            ans[tmp] = 1;

            head[tmp] = 0;


        }
        sort(qu, qu + tail, cmp);
        root = qu[0];

        build();

        dp(root, -1);
        dp1(root, -1);

        solve(root, -1);

        printf("%d", ans[qu0[0]]);
        for (int i = 1; i < tail; i++)
        {
            printf(" %d", ans[qu0[i]]);
        }
        printf("\n");


        q--;


    }



    return 0;
}