【每日一题Day354】LC2316统计无向图中无法互相到达点对数 | 并查集

class Solution {
    public long countPairs(int n, int[][] edges) {
        // 并查集记录每个点的父节点，父节点相同的节点可以相互到达；与父节点不同的节点，不可以互相到达
        // 计算数量，父节点为i的节点数量为a，那么这个节点与n-a个节点不相邻，统计数量
        // 统计结果，由于重复计算，最终结果需要除以2
        UnionFind union = new UnionFind(n);
        for (int[] edge : edges){
            union.union(edge[0], edge[1]);
        }
        long res = 0L;
        for (int i = 0; i < n; i++){
            res += n - union.getSize(i);
        }
        return res / 2;
    }
}
// 并查集类
class UnionFind {
    private int[] p;// 父节点
    private int[] size;// 每个节点的子节点个数（包括自身）
  // 初始化
    public UnionFind(int n) {
        p = new int[n];
        size = new int[n];
        for (int i = 0; i < n; ++i) {
            p[i] = i;
            size[i] = 1;
        }
    }
  // 找到节点x的根
    public int find(int x) {
        if (p[x] != x) {
            p[x] = find(p[x]);
        }
        return p[x];
    }
  // 在并查集中加入a->b的根
    public void union(int a, int b) {
        int pa = find(a), pb = find(b);
        if (pa != pb) {
            if (size[pa] > size[pb]) {
                p[pb] = pa;
                size[pa] += size[pb];
            } else {
                p[pa] = pb;
                size[pb] += size[pa];
            }
        }
    }
    public int getSize(int a){
        int father = find(a);
        return size[father];
    }
}