MalwareSpread.py

'''
Accepted on leetcode
T = O(m*n)
S= O(n)
'''

class Solution:
    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:
        n = len(graph)# finding the length of the graph
        self.colors= [-1 for i in range(n)]
        c = 0  
        for i in range(n):
                # call dfs - to build the  colors array
                self.dfs(graph,i, c)
                c += 1
        
        # initializing the group array to 0
        group = [0 for i in range(c)]
        
        #Finding the number of elemnts belonging to same color
        for i in range(len(self.colors)):
            group[self.colors[i]]+=1

        #Initializing the initial color array
        initcolors = [0 for i in range(c)]
        
        # finding the elements in the initial array belonging to a particular color 
        for i in initial:
            initcolors[self.colors[i]]+=1
        # print(initcolors)
        
        result = float('inf')
        
        for node in initial:
            col = self.colors[node] #Taking the color of a node
            count = initcolors[col] # checking how many elements are there of same color in the initcolor array
            if(count == 1):
                if(result == float('inf')):
                    result = node
                #Checking if the affected group size is greater than the other affected group size
                elif(group[col] > group[self.colors[result]]):
                    result = node
                # if the group with once affected node have the same size the  return the group with minimum node value
                elif(group[col] == group[self.colors[result]] and node < result):
                    result = node
        if(result == float('inf')):
            for node in initial:
                result = min(result,node)
        return result
                    
        


    def dfs(self, graph,node,c):
        
        if self.colors[node] != -1:
            return
        self.colors[node] = c
        for x in range(len(graph)):
            if graph[node][x] == 1:
                self.dfs(graph,x, c)