'''
URL: https://leetcode.com/problems/minimize-malware-spread
Time complexity: O(V+E)
Space complexity: O(V+E)
'''

from collections import deque

class Solution:
    def minMalwareSpread(self, graph, initial):
        """
        :type graph: List[List[int]]
        :type initial: List[int]
        :rtype: int
        """
        if len(initial) == 0:
            return -1
        if len(initial) == 1:
            return initial[0]

        inf_nodes = set(initial)
        visited = set()


        smallest_inf_connected_node = float('inf')
        total_nodes = -1

        i = 0
        while i < len(initial):

            if initial[i] not in visited:
                pot_total_nodes, smallest_inf_node = self._bfs(graph, visited, inf_nodes, initial[i])
                if pot_total_nodes > total_nodes or (pot_total_nodes == total_nodes and smallest_inf_node < smallest_inf_connected_node) :
                    smallest_inf_connected_node = smallest_inf_node
                    total_nodes = pot_total_nodes

            i += 1

        return smallest_inf_connected_node

    def _bfs(self, graph, visited, inf_nodes, source_node):
        queue = deque([source_node])
        visited.add(source_node)
        smallest_inf_node = source_node
        connected_inf_nodes = 1
        total_nodes = 1

        while queue:
            curr_node = queue.popleft()

            for neighbour_node, is_connected in enumerate(graph[curr_node]):
                if not is_connected or neighbour_node in visited:
                    continue

                if neighbour_node in inf_nodes:
                    smallest_inf_node = min(smallest_inf_node, neighbour_node)
                    connected_inf_nodes += 1

                total_nodes += 1

                queue.append(neighbour_node)
                visited.add(neighbour_node)

        return total_nodes, smallest_inf_node
