Minimize Malware Spread II LeetCode solution in Python

1
from collections import defaultdict
2
from collections import deque
3

4

5
class Solution:
6
    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:
7
        n = len(graph)
8
        G = defaultdict(list)
9
        for v in range(n):
10
            for w in range(n):
11
                if graph[v][w]:
12
                    G[v].append(w)
13

14
        affect = defaultdict(list)
15
        for s in initial:
16
            visited = set([v for v in initial if v != s])
17
            que = deque([s])
18
            while que:
19
                v = que.popleft()
20
                if v in visited:
21
                    continue
22
                visited.add(v)
23
                affect[v].append(s)
24
                for w in G[v]:
25
                    que.append(w)
26
        res = [0] * n
27
        for v in affect:
28
            if len(affect[v]) == 1:
29
                res[affect[v].pop()] += 1
30
        if not max(res):
31
            return min(initial)
32
        return res.index(max(res))