from collections import defaultdict
from union_find import *

def min_malware_spread(graph, initial):
    # Stores the length of the graph
    length = len(graph)
    # Calls UnionFind constructor
    union_find = UnionFind(length)
    # Find all the connected components of the graph
    for x in range(length):
        for y in range(length):
            if graph[x][y]:
                union_find.union(x, y)

    infected = defaultdict(int)

    # Count the number of initial infected nodes each connected component has
    for x in initial:
        infected[union_find.find(x)] += 1

    maximum_size, candidate_node = 0, min(initial)
    # Count all the infected nodes each connected component has
    for i in initial:
        infection_count = infected[union_find.find(i)]
        component_size = union_find.rank[union_find.find(i)]

        if infection_count != 1:
            continue
        # Return the candidate node from largest length connected component
        if component_size > maximum_size:
            maximum_size = component_size
            candidate_node = i
        elif component_size == maximum_size and i < candidate_node:
            candidate_node = i
    return candidate_node

# Driver code
def main():
    graph = [[[1, 1, 1], [1, 1, 1], [1, 1, 1]],
            [[1, 1, 0],[1, 1, 0], [0, 0, 1]],
            [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
            [[1, 1, 1, 0, 1], [1, 1, 1, 0, 1], [1, 1, 1, 0, 1], [1, 1, 1, 0, 1], [1, 1, 1, 0, 1]],
            [[1, 1, 1, 0, 1], [1, 1, 1, 0, 1], [1, 1, 1, 0, 1], [1, 1, 1, 0, 1], [1, 1, 1, 0, 1]]]
    initial = [[1, 2], [0, 1], [0, 2], [2, 3], [3, 4]]
    for i in range(0, len(initial)):
        print(i+1, ".\tgraph = ", graph[i], ", ", "Initial = ", initial[i], sep="")
        print("\n\tThe node which minimizes malware spread is", min_malware_spread(graph[i], initial[i]))
        print("-"*100)


if __name__ == '__main__':
    main()