hdu 5441 travel 离线+带权并查集

  1 #include<cstdio>
  2 #include<cstring>
  3 #include<algorithm>
  4 #include<set>
  5 
  6 using namespace std;
  7 
  8 struct Edge
  9 {
 10     int u,v,w;
 11 };
 12 Edge edge[100000+5];
 13 
 14 struct Query
 15 {
 16     int x,id;
 17 };
 18 Query qry[5005];
 19 
 20 int w[20000+5];
 21 long long print[5005];
 22 int fa[20000+5];
 23 
 24 void init(int n)
 25 {
 26     for(int i=1;i<=n;i++){
 27         fa[i]=i;
 28         w[i]=1;
 29     }
 30 }
 31 
 32 bool cmp1(Edge z,Edge y)
 33 {
 34     return z.w<y.w;
 35 }
 36 bool cmp2(Query z,Query y)
 37 {
 38     return z.x<y.x;
 39 }
 40 
 41 int find_fa(int a)
 42 {
 43     if(fa[a]==a)
 44         return a;
 45     else{
 46         fa[a]=find_fa(fa[a]);
 47         w[fa[a]]+=w[a];
 48         w[a]=0;
 49         return fa[a];
 50     }
 51 }
 52 
 53 void solve(int ,int ,int );
 54 
 55 int main()
 56 {
 57     int test;
 58     scanf("%d",&test);
 59     while(test--){
 60         int n,m,Q;
 61         scanf("%d %d %d",&n,&m,&Q);
 62         init(n);
 63         for(int i=1;i<=m;i++){
 64             scanf("%d %d %d",&edge[i].u,&edge[i].v,&edge[i].w);
 65         }
 66         for(int i=1;i<=Q;i++){
 67             scanf("%d",&qry[i].x);
 68             qry[i].id=i;
 69         }
 70         solve(n,m,Q);
 71     }
 72     return 0;
 73 }
 74 
 75 void solve(int n,int m,int Q)
 76 {
 77     sort(edge+1,edge+m+1,cmp1);
 78     sort(qry+1,qry+Q+1,cmp2);
 79 
 80     int cnt=1;
 81     long long sum=0;
 82     for(int i=1;i<=Q;i++){
 83         while(cnt<=m&&edge[cnt].w<=qry[i].x){
 84             int fau=find_fa(edge[cnt].u);
 85             int fav=find_fa(edge[cnt].v);
 86             if(fau!=fav){
 87                 sum+=w[fau]*w[fav]*2;
 88                 fa[fau]=fav;
 89                 w[fav]+=w[fau];
 90                 w[fau]=0;
 91             }
 92             cnt++;
 93         }
 94         print[qry[i].id]=sum;
 95     }
 96     
 97     for(int i=1;i<=Q;i++){
 98         printf("%I64d\n",print[i]);
 99     }
100     
101     return ;
102 }