HDU-5441-离线化并查集

#include <iostream>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <vector>
#include <algorithm>
#include <queue>
#include <set>
#include <map>
#include <stack>
#include <iomanip>
#include <string>

using namespace std;

int t,n,m,q;
int f[20005],s[20005];
struct Road{
    int u,v,cost;
    bool operator<(const Road &a)const {
        return cost<a.cost;
    }
}road[100005];
struct Query{
    int dis,id,ans;
}query[5005];
void ini(){
    for(int i=0;i<=n;i++){
        f[i]=i;
        s[i]=1;//s数组存储和i同属一个并查集的节点数
    }
}
bool cmp1(Query a,Query b){
    return a.dis<b.dis;
}
bool cmp2(Query a,Query b){
    return a.id<b.id;
}
int Find(int x){
    if(x==f[x]) return x;
    else return f[x]=Find(f[x]);
}
int main()
{
    scanf("%d",&t);
    while(t--){
        scanf("%d%d%d",&n,&m,&q);
        ini();
        for(int i=0;i<m;i++)
            scanf("%d%d%d",&road[i].u,&road[i].v,&road[i].cost);
        sort(road,road+m);//先将道路根据值从小到大排序
        for(int i=0;i<q;i++){
            scanf("%d",&query[i].dis);
            query[i].id=i;
        }
        sort(query,query+q,cmp1);//将查询条件递增排序，从而使得离线处理并查集
        int j=0,sum=0;
        for(int i=0;i<q;i++){
            for(;j<m;j++){
                if(road[j].cost>query[i].dis) break;//不满足要求就开始处理下一个查询
                int u=Find(road[j].u),v=Find(road[j].v);
                if(u!=v){//合并两个并查集
                    sum+=s[u]*s[v]*2;//互相可达，那么就是两个并查集点的乘积，根据题意双向算两条
                    f[v]=u;
                    s[u]+=s[v];//要记得更新根节点有多少个节点
                }
            }
            query[i].ans=sum;
        }
        sort(query,query+q,cmp2);
        for(int i=0;i<q;i++)
            printf("%d\n",query[i].ans);
    }
    return 0;
}