2316. 统计无向图中无法互相到达点对数（并查集模板）

Python ml

于 2024-08-14 12:55:15 发布

阅读量119

点赞数 3

分类专栏： LeetCode 文章标签： python 算法开发语言

本文链接：https://blog.csdn.net/m0_48711099/article/details/141188127

版权

LeetCode 专栏收录该内容

176 篇文章 0 订阅

订阅专栏

class UnionFindSet:
    def __init__(self, n: int) -> None:
        self.parent = [i for i in range(n)]
        self.rank = [0] * n

    def find(self, x: int) -> int:
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x: int, y: int) -> None:
        xroot, yroot = self.find(x), self.find(y)
        if xroot != yroot:
            if self.rank[xroot] < self.rank[yroot]:
                xroot, yroot = yroot, xroot
            self.parent[yroot] = xroot
            if self.rank[xroot] == self.rank[yroot]:
                self.rank[xroot] += 1

class Solution:
    def countPairs(self, n: int, edges: List[List[int]]) -> int:
        
        graph = defaultdict(list)
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)
            
        ufs = UnionFindSet(n)
        for u, vset in graph.items():
            for v in vset:
                ufs.union(u, v)
                
        counter = defaultdict(int)
        for u in range(n):
            counter[ufs.find(u)] += 1
            
        res = 0
        for cnt in counter.values():
            res += cnt * (n - cnt)

        res //= 2

        return res