minimizeMalwareSpread.py

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

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

class Solution:
    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:
        # color
        n = len(graph)
        colors = {}
        color = 0
        
        def dfs(node, clor):
            colors[node] = clor
            for i in range(len(graph[node])):
                if graph[node][i] and i not in colors:
                    dfs(i, clor)
        
        for node in range(n):
            if node not in colors:
                dfs(node, color)
                color += 1
        
        colorcnt = collections.Counter(colors.values())
        
        # initial
        inicnt = collections.Counter()
        for node in initial:
            inicnt[colors[node]] += 1
        
        res = float('inf')
        for node in initial:
            c = colors[node]
            if inicnt[c] == 1:
                if res == float('inf'):
                    res = node
                elif colorcnt[c] > colorcnt[colors[res]]:
                    res = node
                elif colorcnt[c] == colorcnt[colors[res]] and node < res:
                    res = node
        
        return res if res < float('inf') else min(initial)