Minimize Malware Spread LeetCode solution in Python

1
# Runtime: 1782 ms (Top 87.91%) | Memory: 16.6 MB (Top 91.21%)
2
class Solution:
3
    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:
4
        n = len(graph)
5
        uf = UnionFind(n)
6
        for i in range(n):
7
            for j in range(n):
8
                if graph[i][j]:
9
                    uf.union(i, j)
10

11
        max_remove = 0
12

13
        initial.sort()
14
        min_index = initial[0]
15
        for i in range(len(initial)):
16
            cur_remove = 0
17
            linked = False
18
            for j in range(len(initial)):
19
                if i == j:
20
                    continue
21
                if uf.find(initial[i]) == uf.find(initial[j]):
22
                    linked = True
23
            if not linked:
24
                cur_remove = uf.rank[uf.find(initial[i])]
25
                if max_remove < cur_remove:
26
                    max_remove = cur_remove
27
                    min_index = initial[i]
28

29
        return min_index
30

31

32
class UnionFind:
33
    def __init__(self, size):
34
        self.parent = {}
35
        self.rank = {}
36
        for i in range(size):
37
            self.parent[i] = i
38
            self.rank[i] = 1
39

40
    def find(self, x):
41
        if x != self.parent[x]:
42
            x = self.find(self.parent[x])
43
        return x
44

45
    def union(self, x, y):
46
        px, py = self.find(x), self.find(y)
47
        if px == py:
48
            return
49
        if self.rank[px] >= self.rank[py]:
50
            self.parent[py] = px
51
            self.rank[px] += self.rank[py]
52
        else:
53
            self.parent[px] = py
54
            self.rank[py] += self.rank[px]