hdu 4750 Count The Pairs(并查集)

题目链接：http://acm.hdu.edu.cn/showproblem.php?pid=4750

代码：

#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
using namespace std;

const int maxn = 10005;
const int maxm = 500500;

struct Edge
{
    int u,v,w;
    Edge(int u=0,int v=0,int w=0): u(u), v(v), w(w) {}
    bool operator < (const Edge& rhs) const
    {
        return w < rhs.w;
    }
}edges[maxm];
struct Query
{
    int id;
    int num;
    bool operator < (const Query& rhs) const
    {
        return num < rhs.num;
    }
}Q[maxn*10];

int pa[maxn];
int counts[maxn];
int ans[maxn*10];

int find(int x)
{
    return x == pa[x] ? x : pa[x] = find(pa[x]);
}


int main()
{
    //freopen("E:\acm\input.txt","r",stdin);
    int n,m,p;
    while(cin>>n>>m)
    {
        for(int i=0; i<n; i++)
        {
            pa[i] = i;
            counts[i] = 1;
        }
        for(int i=0; i<m; i++)
        {
            int a,b,c;
            scanf("%d %d %d",&a,&b,&c);
            edges[i] = Edge(a,b,c);
        }
        sort(edges,edges+m);

        scanf("%d",&p);
        for(int i=0; i<p; i++)
        {
            scanf("%d",&Q[i].num);
            Q[i].id = i;
        }
        sort(Q,Q+p);

        int cnt = 0;
        int sum = 0;
        int tot = n*(n-1);   //这么做前提是图是连通的，只能样例给的是连通的

        for(int i=0; i<p; i++)
        {
            while(cnt<m && edges[cnt].w < Q[i].num)
            {
                int u_fa = find(edges[cnt].u);
                int v_fa = find(edges[cnt].v);
                if(u_fa != v_fa)
                {
                    sum = sum + counts[u_fa]*counts[v_fa]*2;
                    pa[v_fa] = u_fa;
                    counts[u_fa] += counts[v_fa];
                }
                cnt++;
            }

            ans[Q[i].id] = tot - sum;
        }
        for(int i=0; i<p; i++)
            printf("%d
",ans[i]);
    }
}