并查集 HDOJ 5441 Travel

/************************************************
 * Author        :Running_Time
 * Created Time  :2015/9/13 星期日 17:55:11
 * File Name     :E.cpp
 ************************************************/

#include <cstdio>
#include <algorithm>
#include <iostream>
#include <sstream>
#include <cstring>
#include <cmath>
#include <string>
#include <vector>
#include <queue>
#include <deque>
#include <stack>
#include <list>
#include <map>
#include <set>
#include <bitset>
#include <cstdlib>
#include <ctime>
using namespace std;

#define lson l, mid, rt << 1
#define rson mid + 1, r, rt << 1 | 1
typedef long long ll;
const int N = 2e4 + 10;
const int M = 1e5 + 10;
const int Q = 5e3 + 10;
const int INF = 0x3f3f3f3f;
const int MOD = 1e9 + 7;
struct UF   {
    int rt[N], rk[N];
    void init(void) {
        memset (rt, -1, sizeof (rt));
        memset (rk, 0, sizeof (rk));
    }
    int Find(int x) {
        return rt[x] == -1 ? x : rt[x] = Find (rt[x]);
    }
    void Union(int x, int y)    {
        x = Find (x), y = Find (y);
        if (x == y) return ;
        if (rk[x] > rk[y])  {
            rt[y] = x;  rk[x] += rk[y] + 1;
        }
        else    {
            rt[x] = y;  rk[y] += rk[x] + 1;
        }
    }
    bool same(int x, int y) {
        return Find (x) == Find (y);
    }
}uf;
struct Edge {
    int u, v, w;
    bool operator < (const Edge &r) const   {
        return w < r.w;
    }
}edge[M];
struct Query    {
    int x, id;
    bool operator < (const Query &r) const  {
        return x < r.x;
    }
}q[Q];
ll ans[Q];

int main(void)    {
    int T;  scanf ("%d", &T);
    while (T--) {
        int n, m, k;
        scanf ("%d%d%d", &n, &m, &k);
        for (int i=1; i<=m; ++i)    {
            scanf ("%d%d%d", &edge[i].u, &edge[i].v, &edge[i].w);
        }
        sort (edge+1, edge+1+m);
        for (int i=1; i<=k; ++i)    {
            scanf ("%d", &q[i].x);  q[i].id = i;
        }
        sort (q+1, q+1+k);

        memset (ans, 0, sizeof (ans));
        uf.init (); int j = 1;  ll sum = 0;
        for (int i=1; i<=k; ++i)    {
            while (j <= m && q[i].x >= edge[j].w)   {
                int u = edge[j].u, v = edge[j].v;
                if (uf.same (u, v)) {
                    ++j;    continue;
                }
                u = uf.Find (u), v = uf.Find (v);
                sum += (uf.rk[u] + 1) * 1ll * (uf.rk[v] + 1);
                uf.Union (u, v);    ++j;
            }
            ans[q[i].id] = sum;
        }
        for (int i=1; i<=k; ++i)    {
            printf ("%I64d\n", ans[i] * 2);
        }
    }

    return 0;
}