HDU - 5441- Travel (并查集)

/*
按 x从小到大处理：
按边长从小到大处理:
可达的一条边的两个点视作一个并查集
合并即可
*/

#include <bits/stdc++.h>
using namespace std;
const int maxn = 20000 + 10;
int father[maxn],num[maxn];
int n,m,q,p,t;
int a[maxn];
struct node
{
    int x,index;
    bool operator < (const node &a)const
    {
        return x < a.x;
    }
}query[maxn];
struct Edge
{
    int u,v,w;
    bool operator < (const Edge & e)const
    {
        return w < e.w;
    }
}e[maxn*5];
void init()
{
    for(int i = 1;i <= n; i++)
    {
        num[i] = 1;
        father[i] = i;
    }
}
int find_ ( int x )
{
    return father[x] == x? x : father[x] = find_ ( father[x] );
}
void union_(int x,int y)
{
    x = find_(x);
    y = find_(y);
    if(y < x) swap(x,y);
    father[y] = x;
    num[x] += num[y];
}
int main()
{
    scanf("%d",&t);
    while(t--)
    {
        int ans = 0;
        scanf("%d%d%d",&n,&m,&q);
        init();
        for(int i = 0; i< m;i++)
        {
            scanf("%d%d%d",&e[i].u,&e[i].v,&e[i].w);
        }
        sort(e,e+m);
        for(int i = 0; i < q; i++)
        {
            query[i].index = i;
            scanf("%d",&query[i].x);
        }
        sort(query,query + q);
        int j = 0;
        for(int i = 0 ; i < q; i++)
        {
            while(j < m && e[j].w <= query[i].x)
            {
                int u = find_(e[j].u);
                int v = find_(e[j].v);
                j++;
                if(u == v) continue;
                ans += (num[u]+num[v])*(num[u]+num[v]-1)-num[u]*(num[u]-1) - num[v]*(num[v]-1);
                union_ ( u , v );
            }
            a[query[i].index] = ans;
        }

        for(int i = 0 ; i < q;i++)
            printf("%d\n",a[i]);
    }
    return 0;
}