SMMR-198: Pakiny & C-Distance

#include <bits/stdc++.h>

using namespace std;
using lli = long long;
const int N = 1e5;
using pi = pair <int, int>;
vector <pi> g[N+10];
int dis[N+10], qs[N+10], d[N+10];
bool vs[N+10], node[N+10];
int n, m, c;

int Cnt(int u){
    if(vs[u]) return 0;
    vs[u] = true;
    int cnt = 1;
    for(auto vw: g[u]){
        int v = vw.first;
        int w = vw.second;
        if(w >= c) continue;
        cnt += Cnt(v);
    }
    return cnt;
}

bool cpn[N+10];
void Cpn(int u){
    if(cpn[u]) return;
    cpn[u] = true;
    for(auto vw: g[u]){
        int v = vw.first;
        int w = vw.second;
        if(cpn[v]) continue;
        Cpn(v);
    }
}

int main(){

    scanf("%d%d%d", &n, &m, &c);

    for(int i=1;i<=m;i++){
        int u, v, w;
        scanf("%d%d%d", &u, &v, &w);
        if(w > c) continue;
        g[u].push_back({v, w});
        g[v].push_back({u, w});
        if(w == c) node[v] = node[u] = true;
    }

    for(int u=1;u<=n;u++){
        if(node[u]) dis[u] = Cnt(u);
    }
    lli ans = 0;
    for(int u=1;u<=n;u++){
        if(node[u]){
            Cpn(u);
            int sz = 0;
            for(int i=1;i<=n;i++){
                if(cpn[i] and node[i]){
                    d[sz] = dis[i];
                    sz ++;
                }
            }

            qs[sz-1] = d[sz-1];
            for(int i=sz-2;i>=0;i--)
                qs[i] = qs[i+1] + d[i];
            for(int i=0;i<sz-1;i++)
                ans += (lli) d[i] * qs[i+1];
            for(int i=1;i<=n;i++)
                if(cpn[i]) cpn[i] = node[i] = false;
        }
    }

    printf("%lld\n", ans);

    return 0;
}
/**

10 11 4
1 2 4
3 2 4
4 2 3
5 2 5
4 6 5
5 6 4
5 9 2
9 10 1
8 9 4
7 8 7
4 7 6

*/