minimizeMalwareSpreadII.py

#! /usr/bin/env python3
# -*- coding: utf-8 -*-

# Source: https://leetcode.com/problems/minimize-malware-spread-ii/
# Author: Miao Zhang
# Date:   2021-03-26

class Solution:
    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:
        n = len(graph)
        clean = set(range(n)) - set(initial)
        def dfs(u, seen):
            for i in range(len(graph[u])):
                if graph[u][i] and i in clean and i not in seen:
                    seen.add(i)
                    dfs(i, seen)   
        infected_by = {v: [] for v in clean}
        for u in initial:
            seen = set()
            dfs(u, seen)
            for v in seen:
                infected_by[v].append(u)
        
        contribution = collections.Counter()
        for k, v in infected_by.items():
            if len(v) == 1:
                contribution[v[0]] += 1
        
        best = (-1, min(initial))
        for u, score in contribution.items():
            if score > best[0] or score == best[0] and u < best[1]:
                best = score, u
        return best[1]