Leetcode 964 Solution

class Solution:
    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:
        initial_set = set(initial)
        neighbors = collections.defaultdict(list)

        n = len(graph)
        for i in range(n):
            for j in range(n):
                if graph[i][j] == 1:
                    neighbors[i].append(j)

        all_impacted_nodes = {}
        for node in initial:
            initial_node = node
            impacted_nodes = []

            q = collections.deque()
            q.append(node)
            visited = {node}
            while q:
                node = q.popleft()
                impacted_nodes.append(node)

                for neigh_node in neighbors[node]:
                    if neigh_node in visited or neigh_node in initial_set:
                        continue

                    visited.add(neigh_node)
                    q.append(neigh_node)

            all_impacted_nodes[initial_node] = impacted_nodes

        node_taint_count = collections.defaultdict(int)
        for initial_node, impacted_nodes in all_impacted_nodes.items():
            for impacted_node in impacted_nodes:
                node_taint_count[impacted_node] += 1

        candidates = {}
        opt_cnt = 0
        for initial_node, impacted_nodes in all_impacted_nodes.items():
            cnt = 0
            for impacted_node in impacted_nodes:
                if node_taint_count[impacted_node] == 1:
                    cnt += 1
            candidates[initial_node] = cnt
            opt_cnt = max(opt_cnt, cnt)

        for initial_node, cnt in sorted(candidates.items()):
            if cnt == opt_cnt:
                return initial_node