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

· · 题解

题意

给定 n 个点的树与 q 次询问，每次询问给定一个大小为 k 的点集 S，求 S 中有多少点对满足两点简单路径上仅包括 S 中的点。

题解

朴素做法是 O(qn) 的，很劣。注意到 \sum k \le 10^6，考虑从它入手。对于一个点 x \in S，若其父节点 fa_x 满足 fa_x \in S，则 (x，fa_x) 是一个满足条件的点对。由此推广，考虑维护一个并查集，枚举 S 中的点 x，若 fa_x \in S，则将 x 与 fa_x 连边。求出所有联通块的数量 m 和大小 siz，则答案为 \sum ^{m} _{i=1} \frac{siz_i \times (siz_i-1)}{2}。

复杂度 O(\sum k)，我的写法加上快读就进最优解前十了，还挺快。

代码

#include<bits/stdc++.h>
using namespace std;
#define int long long
const int N=3e5+5;
int n，q，k，a[N]，p[N];
bool v[N];
vector<int>e[N];
int ans;
void dfs(int x，int f)
{
    p[x]=f;
    for(auto y:e[x])
        if(y!=f)
            dfs(y，x);
}
int fa[N]，siz[N];
int find(int x)
{return fa[x]=fa[x]==x?x:find(fa[x]);}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
    cin>>n;
    for(int i=1，x，y;i<n;i++)
        cin>>x>>y，
        e[x].push_back(y)，
        e[y].push_back(x);
    dfs(1，0);
    cin>>q;
    for(;q;q--)
    {
        cin>>k;ans=0;
        for(int i=1;i<=k;i++)
            cin>>a[i];
        for(int i=1;i<=k;i++)
            v[a[i]]=1，fa[a[i]]=a[i];
        for(int i=1;i<=k;i++)
            if(v[a[i]]&&v[p[a[i]]])
                fa[a[i]]=find(p[a[i]]);
        for(int i=1;i<=k;i++)
            siz[find(a[i])]++;
        for(int i=1;i<=k;i++)
            ans+=siz[a[i]]*(siz[a[i]]-1)/2;
        for(int i=1;i<=k;i++)
            v[a[i]]=0，siz[a[i]]=0;
        cout<<ans<<'\n';
    }
    return 0;
}