HDU 5441 Travel

HDU 5441 Travel

#include<algorithm>
#include<iostream>
#include<cstring>
#include<cstdio>
#include<string>
#include<vector>
#include<cmath>
#define maxn 100005
#define INF 0x3fffffff
#define ll __int64
#define mod 10000007
#define N 20005
#define M 100005
using namespace std;
int s[N];
int que[M];
struct node{
    int a,b;
    int val;
}p[M];
ll ans[M];
int Find(int num){
    if(s[num]<0) return num;
    return s[num]=Find(s[num]);
}
bool cmp(node a,node b){
    return a.val<b.val;
}
int main(){
    int T;
    scanf("%d",&T);
    int n,m,q;
    while(T--){
        memset(s,-1,sizeof(s));
        scanf("%d%d%d",&n,&m,&q);
        for(int i=0;i<m;i++)
            scanf("%d%d%d",&p[i].a,&p[i].b,&p[i].val);
        sort(p,p+m,cmp);
        int E=0;
        for(int i=0;i<q;i++){
            scanf("%d",&que[i]);
            E=max(E,que[i]);
        }
        /***
        Init
        */
        ll sum=0,A=0;
        for(int i=0;i<=E&&A<m;i++){
            while(p[A].val<=i&&A<m){
                int xx=Find(p[A].a);
                int yy=Find(p[A].b);
                if(xx!=yy){
                    ll nu1=-s[xx];
                    ll nu2=-s[yy];
                    s[xx]+=s[yy];
                    s[yy]=xx;
                    sum-=nu1*(nu1-1);
                    sum-=nu2*(nu2-1);
                    sum+=(nu1+nu2)*(nu1+nu2-1);
                }
                A++;
            }
            ans[i]=sum;
        }
        for(int i=0;i<q;i++)
            printf("%I64d\n",ans[que[i]]);
    }
    return 0;
}