HDU 4750 Count The Pairs

#include <stdio.h>
#include <string.h>
#include <algorithm>
#define maxn 500010
using namespace std;
int par[10010], rank1[10010], a[maxn], b[maxn];
long long int sum[maxn];
void init(int n){
for(int i = 0;i < n;i++){
par[i] = -1;
rank1[i] = 1;
}
}
int find(int x){
if(par[x] == -1) return x;
else return par[x] = find(par[x]);
}
void unite(int x, int y){
x = find(x);
y = find(y);
if(x != y)
{
par[x] = y;
rank1[y] += rank1[x];
}
}

struct node
{
int u, v, l;
}edge[maxn];
bool cmp(node a, node b){
return a.l < b.l;
}
int main()
{
int i, n, m, p, tp, id;
while(scanf("%d %d", &n, &m) != EOF)
{
for(i = 0;i < m;i++) scanf("%d %d %d", &edge[i].u, &edge[i].v, &edge[i].l);
init(n);
sort(edge, edge + m, cmp);

for(i = 0;i < m;i++)
{
b[i] = edge[i].l;
int u = edge[i].u;
int v = edge[i].v;
if(find(u) != find(v))
{
a[i] = 2*rank1[find(u)]*rank1[find(v)];
unite(u, v);
}
else a[i] = 0;
}
sum[m] = 0;
for(i = m - 1;i >= 0;i--) sum[i] = sum[i + 1] + a[i];
scanf("%d", &p);
for(i = 0;i < p;i++)
{
scanf("%d", &tp);
id = lower_bound(b, b + m, tp) - b;
if(id >= m) printf("0\n");
else printf("%lld\n", sum[id]);
}
}
return 0;
}