hdu 5441 Travel 离线操作+并查集

#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
const int MAXN = 5005;
const int INF  = 0x3f3f3f3f;
int pre[MAXN<<2];
int M,N;

struct Node{
    int a,b,val;
}e[MAXN<<5];
bool cmp(Node a,Node b)
{
    return a.val<b.val;
}
struct Node2{
    int bre,id;
}s[MAXN];
bool cmp2(Node2 a,Node2 b)
{
    return a.bre<b.bre;
}
int w[MAXN<<2];
int findSet(int x)
{
    return (x==pre[x])?pre[x]:(pre[x] = findSet(pre[x]));
}
int main()
{
    int T,Q;
    scanf("%d",&T);
    while(T--)
    {
        int ans[MAXN];
        scanf("%d%d%d",&N,&M,&Q);
        for(int i=1;i<=N;i++)
            pre[i] = i,w[i] = 1;
        for(int i=1;i<=M;i++)
            scanf("%d%d%d",&e[i].a,&e[i].b,&e[i].val);
        sort(e+1,e+M+1,cmp);
        for(int i=1;i<=Q;i++)
            scanf("%d",&s[i].bre),s[i].id = i;
        sort(s+1,s+Q+1,cmp2);
        int amount = 0,j = 1;
        for(int i=1;i<=Q;i++)
        {
            for(;j<M&&e[j].val<=s[i].bre;j++)
            {
                int fx = findSet(e[j].a),fy = findSet(e[j].b);
                if(fx!=fy)
                {
                    pre[fx] = fy;
                    amount -= (w[fx]*(w[fx]-1) + w[fy]*(w[fy]-1));
                    w[fy] += w[fx];
                    amount += w[fy]*(w[fy]-1);
                }
            }
            ans[s[i].id] = amount;
        }
        for(int i=1;i<=Q;i++)
            printf("%d\n",ans[i]);
    }
    return 0;
}