City - 并查集路径压缩与按秩合并 datastructure

 

 

 

 

 

 

#include <bits/stdc++.h>

using i64 = long long;
#define dbg(x) std::cerr << #x << "=" << x << "\n"

struct node{
    int u,v,w;
};

int fa[100005], rank[100005],sizes[100005];
int find(int x){
    return fa[x] == x ? x : (fa[x] = find(fa[x]));
}
inline i64 C2(i64 k){
    return k * (k - 1) / 2;
}
//并查集的路径压缩和按秩合并都要用
void merge(int u, int v, std::vector<i64>& ans){
    u = find(u);
    v = find(v);
    if(rank[u] > rank[v]){
        fa[v] = u;
        int t1 = sizes[u] == 1 ? 0 : sizes[u];
        int t2 = sizes[v] == 1 ? 0 : sizes[v];
        if(u != v) sizes[u] += sizes[v];
        i64 res = C2(sizes[u]);
        if(sizes[u] == t1)//如果包含的点集数没有变化就要把变化值res变成0
            res = 0;
        else//若有变化，就需要减去之前已经形成的对数
            res -= C2(t1) + C2(t2); 
        ans.push_back(ans.back() + res);
    }else{
        fa[u] = v;
        int t1 = sizes[u] == 1 ? 0 : sizes[u];
        int t2 = sizes[v] == 1 ? 0 : sizes[v];
        if(v != u) sizes[v] += sizes[u];
        if(rank[u] == rank[v]) rank[v] += 1;
        i64 res = C2(sizes[v]);
        if(sizes[v] == t2)
            res = 0;
        else 
            res -= C2(t1) + C2(t2);
        ans.push_back(ans.back() + res);
    }
}
void solve(){
    int n,m,Q;
    std::cin >> n >> m >> Q;
    for(int i = 1; i <= n; i++){
        fa[i] = i,sizes[i] = 1, rank[i] = 0;    
    }
    std::vector<node> adj(m + 1);
    for(int i = 1; i <= m; i++){
        int u,v,w;
        std::cin >> u >> v >> w;
        if(u == v) continue;
        adj[i] = (node){u, v, w};
    }
    std::sort(adj.begin() + 1, adj.end(),[&](node i, node j){
        return i.w > j.w;
    });
    std::vector<i64> ans(1, 0);
    for(int i = 1; i <= m; i++){
        int u = adj[i].u, v = adj[i].v;
        merge(u, v, ans);
    }
    while(Q--){
        int p;
        std::cin >> p;
        int l = 0, r = m;
        while(l < r){
            int mid = (l + r + 1) >> 1;
            if(p <= adj[mid].w)
                l = mid;
            else
                r = mid - 1;
        }
        std::cout << ans[l] << "\n";
    }
}
int main(){
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    int t;
    std::cin >> t;
    while(t--){
        solve();
    }
    return 0;
}