text: # Time:  O(n)
  # Space: O(n)
  
  import collections
  
  
  class UnionFind(object):
      def __init__(self, n):
          self.set = range(n)
  
      def find_set(self, x):
          if self.set[x] != x:
              self.set[x] = self.find_set(self.set[x])  # path compression.
          return self.set[x]
  
      def union_set(self, x, y):
          x_root, y_root = map(self.find_set, (x, y))
          if x_root == y_root:
              return False
          self.set[min(x_root, y_root)] = max(x_root, y_root)
          return True
  
  
  class Solution(object):
      def minMalwareSpread(self, graph, initial):
          """
          :type graph: List[List[int]]
          :type initial: List[int]
          :rtype: int
          """
          union_find = UnionFind(len(graph))
          for i in xrange(len(graph)):
              for j in xrange(i+1, len(graph)):
                  if graph[i][j] == 1:
                      union_find.union_set(i, j)
          union_size = collections.Counter(union_find.find_set(i) for i in xrange(len(graph)))
          malware_count = collections.Counter(union_find.find_set(i) for i in initial)
          return min(initial, key=lambda x: [malware_count[union_find.find_set(x)] > 1,
                                             -union_size[union_find.find_set(x)],
                                             x])
  