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;
}