Solution.py

class Solution:
    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:
        n = len(graph)
        p = list(range(n))
        size = [1] * n

        def find(x):
            if p[x] != x:
                p[x] = find(p[x])
            return p[x]

        for i in range(n):
            for j in range(i + 1, n):
                if graph[i][j] == 1:
                    pa, pb = find(i), find(j)
                    if pa == pb:
                        continue
                    p[pa] = pb
                    size[pb] += size[pa]

        mi = float('inf')
        res = initial[0]
        initial.sort()
        for i in range(len(initial)):
            t = 0
            s = set()
            for j in range(len(initial)):
                if i == j:
                    continue
                if find(initial[j]) in s:
                    continue
                s.add(find(initial[j]))
                t += size[find(initial[j])]
            if mi > t:
                mi = t
                res = initial[i]
        return res