Count Unreachable Pairs of Nodes in an Undirected Graph LeetCode solution in Python

1
# Runtime: 2311 ms (Top 90.13%) | Memory: 73.8 MB (Top 86.18%)
2
"""
3
* Make groups of nodes which are connected
4
    eg., edges = [[0,2],[0,5],[2,4],[1,6],[5,4]]
5

6
    0 ---- 2 1 --- 6 3
7
    | |
8
    | |
9
    5 ---- 4
10

11
    groups will be {0: 4, 1: 2, 3: 1},
12
    i.e 4 nodes are present in group0, 2 nodes are present in group1 and 1 node is present in group3
13

14
* Now, we have [4, 2, 1] as no of nodes in each group, we have to multiply each of no. with remaining
15
    ans = (4 * 2 + 4 * 1) + (2 * 1)
16
    but calculating ans this way will give TLE.
17

18
* if we notice, (4 * 2 + 4 * 1) + (2 * 1), we can combine, equation like this,
19
    4 * 2 + (4 + 2) * 1, using this, we can reduce complexity.
20
    so, if we have count of groups array as [a, b, c, d], ans will be,
21
    ans = a * b + (a + b) * c + (a + b + c) * d
22

23
* will use, union for generating groups.
24
* ps, you can modify UnionFind class as per your need. Have implemented full union-find for beginners.
25
"""
26

27

28
class UnionFind:
29
    def __init__(self, size):
30
        self.root = [i for i in range(size)]
31
        self.rank = [1] * size
32

33
    def find(self, x):
34
        if x == self.root[x]:
35
            return x
36
        self.root[x] = self.find(self.root[x])
37
        return self.root[x]
38

39
    def union(self, x, y):
40
        rootX = self.find(x)
41
        rootY = self.find(y)
42
        if rootX != rootY:
43
            if self.rank[rootX] > self.rank[rootY]:
44
                self.root[rootY] = rootX
45
            elif self.rank[rootX] < self.rank[rootY]:
46
                self.root[rootX] = rootY
47
            else:
48
                self.root[rootY] = rootX
49
                self.rank[rootX] += 1
50

51

52
class Solution:
53
    def countPairs(self, n: int, edges: List[List[int]]) -> int:
54
        dsu = UnionFind(n)
55
        for u, v in edges:
56
            dsu.union(u, v)
57
        C = Counter([dsu.find(i) for i in range(n)])
58
        groupCounts = list(C.values())
59
        ans = 0
60
        firstGroupCount = groupCounts[0]
61
        for i in range(1, len(groupCounts)):
62
            ans += firstGroupCount * groupCounts[i]
63
            firstGroupCount += groupCounts[i]
64
        return ans