题解：P12698 [KOI 2022 Round 2] 树与查询

· · 题解

原问题的询问中，点构成若干个连通块，每个连通块中的点都相互可达。因此对于一个大小为 x 的连通块，每个点都能到达剩余 x-1 个点，其贡献为 \dfrac{x(x-1)}{2}。

一个直接的想法是用 DFS 求给定的点中构成了多少个连通块并统计每个连通块的大小，但由于对某个点的 DFS 要遍历其全部出边，最坏情况下时间复杂度是 O(NQ) 的，无法通过。

考虑均摊每条边，对于一棵树，除了根节点以外，每个节点都有且仅有一条连向父节点的边。因此对每个点定它的边就是连接它和父节点的边。

对于每个询问，如果给定的点集中出现了有父子关系的两点，那么它们之间就存在一条边，可以用并查集维护连边的关系。最后统计并查集中每个连通块的大小即可。

#include<bits/stdc++.h>
using namespace std;

#define ll long long
#define up(l,r,i) for(int i=(l);i<=(r);++i)

ll n,q,k,a[250007],ans,f[250007],fa[250007];
vector<int> G[250007];
map<ll,ll> cnt;
map<ll,bool> vis;

ll getfa(ll u){
    return fa[u]==u?u:fa[u]=getfa(fa[u]);
}

void dfs(int u,int _f){
    f[u]=_f;
    for(int v:G[u]){
        if(v==_f)continue;
        dfs(v,u);
    }
}

int main(){
    cin>>n;
    up(2,n,i){
        ll u,v;cin>>u>>v;
        G[u].push_back(v);G[v].push_back(u);
    }
    dfs(1,0);
    cin>>q;
    while(q--){
        vis.clear();cnt.clear();ans=0;
        cin>>k;up(1,k,i)cin>>a[i],vis[a[i]]=1,fa[a[i]]=a[i];
        up(1,k,i){
            if(vis[a[i]]&&vis[f[a[i]]]){
                ll u=a[i],v=f[a[i]];
                u=getfa(u),v=getfa(v);
                fa[u]=v;
            }
        }
        up(1,k,i)++cnt[getfa(a[i])];
        for(auto tmp:cnt){
            ll qwq=tmp.second;
            ans+=(qwq-1)*(qwq)/2;
        }
        cout<<ans<<'\n';
    }
}