uf1.py

class Solution:
    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:
        self.uf = {}
        def union(x, y):
            self.uf[find(y)] = find(x)
        def find(x):
            self.uf.setdefault(x, x)
            if self.uf[x]!=x:
                self.uf[x] = find(self.uf[x])
            return self.uf[x]
        for i in range(len(graph)):
            for j in range(i+1, len(graph)):
                if graph[i][j]:union(i, j)
        cnt = collections.defaultdict(set)
        for i in range(len(graph)):
            if i in self.uf:
                cnt[find(i)].add(i)
        def get_this_uninfected(i):
            if len(cnt[find(i)] & set(initial)) == 1:
                return len(cnt[find(i)])
            return 0
        uninfected, res = -1, 0
        #print(self.uf)
        #print(cnt)
        for i in sorted(initial):
            this_uninfected = get_this_uninfected(i)
            if this_uninfected>uninfected:
                res = i
                uninfected = this_uninfected
        return res