Number of Operations to Make Network Connected LeetCode solution in Python

1
# Runtime: 772 ms (Top 55.10%) | Memory: 34.3 MB (Top 72.08%)
2
class Solution(object):
3
    def __init__(self):
4
        self.parents = []
5
        self.count = []
6

7
    def makeConnected(self, n, connections):
8
        """
9
        :type n: int
10
        :type connections: List[List[int]]
11
        :rtype: int
12
        """
13
        if len(connections) < n - 1:
14
            return -1
15
        self.parents = [i for i in range(n)]
16
        self.count = [1 for _ in range(n)]
17
        for connection in connections:
18
            a, b = connection[0], connection[1]
19
            self.union(a, b)
20
        return len({self.find(i) for i in range(n)}) - 1
21

22
    def find(self, node):
23
        """
24
        :type node: int
25
        :rtype: int
26
        """
27
        while node != self.parents[node]:
28
            node = self.parents[node]
29
        return node
30

31
    def union(self, a, b):
32
        """
33
        :type a: int
34
        :type b: int
35
        :rtype: None
36
        """
37
        a_parent, b_parent = self.find(a), self.find(b)
38
        a_size, b_size = self.count[a_parent], self.count[b_parent]
39

40
        if a_parent != b_parent:
41
            if a_size < b_size:
42
                self.parents[a_parent] = b_parent
43
                self.count[b_parent] += a_size
44
            else:
45
                self.parents[b_parent] = a_parent
46
                self.count[a_parent] += b_size