0924-minimize-malware-spread.py

# time complexity: O(n^2)
# space complexity: O(n)
import collections
from typing import List


class Solution:
    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:
        N = len(graph)
        colors = {}
        c = 0

        def dfs(node, color):
            colors[node] = color
            for nei, adj in enumerate(graph[node]):
                if adj and nei not in colors:
                    dfs(nei, color)
        for node in range(N):
            if node not in colors:
                dfs(node, c)
                c += 1
        size = collections.Counter(colors.values())
        colorCount = collections.Counter()
        for node in initial:
            colorCount[colors[node]] += 1
        ans = float('inf')
        for x in initial:
            c = colors[x]
            if colorCount[c] == 1:
                if ans == float('inf'):
                    ans = x
                elif size[c] > size[colors[ans]]:
                    ans = x
                elif size[c] == size[colors[ans]] and x < ans:
                    ans = x
        return ans if ans < float('inf') else min(initial)

# time complexity: O(n^2)
# space complexity: O(n)
class UnionFind:
    def __init__(self, n):
        self.parents = [num for num in range(n + 1)]
        self.ranks = [1 for _ in range(n)]

    def find(self, num):
        if self.parents[num] != num:
            self.parents[num] = self.find(self.parents[num])
        return self.parents[num]

    def union(self, x, y):
        rootX = self.find(x)
        rootY = self.find(y)

        if rootX == rootY:
            return
        small, big = sorted([rootX, rootY], key=lambda z: self.ranks[z])
        self.parents[small] = big
        self.ranks[big] += self.ranks[small]


class Solution:
    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:
        ROW = len(graph)
        COL = len(graph[0])
        uf = UnionFind(ROW)
        for r in range(ROW):
            for c in range(COL):
                if graph[r][c]:
                    uf.union(r, c)

        infected = collections.defaultdict(int)
        for num in initial:
            infected[uf.find(num)] += 1

        maxSize = 0
        candidateNode = min(initial)

        for num in initial:
            infectionCount = infected[uf.find(num)]
            componentSize = uf.ranks[uf.find(num)]
            if infectionCount != 1:
                continue
            if componentSize > maxSize:
                maxSize = componentSize
                candidateNode = num
            elif componentSize == maxSize and num < candidateNode:
                candidateNode = num

        return candidateNode


graph = [[1, 1, 0], [1, 1, 0], [0, 0, 1]]
initial = [0, 1]
print(Solution().minMalwareSpread(graph, initial))
graph = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
initial = [0, 2]
print(Solution().minMalwareSpread(graph, initial))
graph = [[1, 1, 1], [1, 1, 1], [1, 1, 1]]
initial = [1, 2]
print(Solution().minMalwareSpread(graph, initial))