2316. 统计无向图中无法互相到达点对数

2316. 统计无向图中无法互相到达点对数

我的代码

const int N = 1e5+10;
int p[N], cnt[N];

class Solution {
public:
    int find(int x) {
        if(p[x]!=x) p[x] = find(p[x]);
        return p[x];
    }

    void merge(int a, int b) {
        int ra = find(a), rb = find(b);
        p[rb] = ra;
    }

    long long countPairs(int n, vector<vector<int>>& edges) {
        for(int i = 0; i < n ;i++) p[i] = i;
        memset(cnt, 0, sizeof cnt);

        for(auto e: edges){
            if(e[0] > e[1]) swap(e[0], e[1]);
            merge(e[0], e[1]);
            int ra = find(e[0]);
        }

        for(int i = 0; i < n; i++){
            p[i] = find(i);
            cnt[p[i]]++;
        }

        vector<int> res;
        for(int i = 0; i < n; i++) {
            if(cnt[i]) res.push_back(cnt[i]);
        }
        
        int f[res.size()+10];
        memset(f, 0, sizeof f);

        for(int i = res.size()-1; i>=0; i--) {
            f[i] = f[i+1] + res[i];
        }

        long long ans = 0;
        for(int i = 0; i < res.size(); i++) {
            ans += 1ll * res[i] * f[i+1];
        }

        return ans;
    }
};

汪乐平代码

const int N = 1e5+10;
int p[N], s[N];

class Solution {
public:
    int find(int x){
        if(p[x]!=x) p[x] = find(p[x]);
        return p[x];
    }

    long long countPairs(int n, vector<vector<int>>& edges) {
        long long all = 1ll * n * (n-1) /2;
        for(int i = 0; i<n;i++)p[i] = i, s[i] = 1;
        for(auto e: edges){
            int ra = find(e[0]), rb = find(e[1]);
            if(ra != rb) {
                s[rb] += s[ra];
                p[ra] = find(rb);
            }
        }
        
        for(int i = 0; i< n; i++){
            if(p[i] != i )continue;
            long long inter = 1ll * s[i] * (s[i] - 1) /2;
            all -= inter;
        }
        return all;
    }
};