Problem statement

https://leetcode.com/problems/minimize-malware-spread-ii/

Solution 1

One way to solve this problem is to use Union Find. Let us remove all infected nodes from our graph and denote this graph by $G$. Then we need to do following steps:

  1. Create clean list of not-infected nodes and initial_set set of infected nodes.
  2. Create DSU on $n$ elements and add all connections between not-infected nodes.
  3. Now, for each node, going from infected node to not-infected node we add d1[dsu.find(v)].append(u): connection from representative element of connected component in $G$ to u.
  4. Check for every connected component, how many connections with infected nodes it have: if we have more than one, than we can not save this component. If we have exactly one, we can save it, so we update counter of saved nodes, d2.
  5. So, in this way we did the following operations: for every infected node we found all components in $G$, which connected with it, such that these components connected only with one infected node: this will be total number of saved nodes. Finally, we iterate over d2 and choose node with biggest value, if we have tie, then with smallest index. If d2 is empty, it means, that we can not save anything, so we need to return minimum of indexes in initial.

Complexity

Time complexity is $O(n^2)$ if we use DSU with ranks and path relaxations (in code it is just simplest DSU without any two of them, but it works OK on leetcode), space complexity is $O(n)$.

Code

class DSU:
    def __init__(self, N):
        self.p = list(range(N))
        self.sz = [1] * N

    def find(self, x):
        if self.p[x] != x:
            self.p[x] = self.find(self.p[x])
        return self.p[x]

    def union(self, x, y):
        xr = self.find(x)
        yr = self.find(y)
        self.p[xr] = yr
        if xr != yr: self.sz[yr] += self.sz[xr]

    def size(self, x):
        return self.sz[self.find(x)]

class Solution(object):
    def minMalwareSpread(self, graph, initial):
        n = len(graph)
        initial_set = set(initial)
        clean = [x for x in range(n) if x not in initial_set]

        dsu = DSU(n)
        for u, v in product(clean, clean):
            if graph[u][v]: dsu.union(u, v)
        
        d1, d2 = defaultdict(list), Counter()
        for u, v in product(initial, clean):
            if graph[u][v]: d1[dsu.find(v)].append(u)
        
        for key, val in d1.items():
            if len(val) == 1: d2[val[0]] += dsu.size(key)
                
        if not d2: return min(initial)      
        return max(d2.items(), key = lambda x: (x[1], -x[0]))[0]

Solution 2

There is also solution with the following idea: we define node a are directly infected by node b if node a will be infected by node b without through any other infected node. Then we can run dfs/bfs from every infected node and not allow it to go on other infected nodes. In the end we find which one from initial nodes appears the most and choose it.

Complexity

Time complexity will be $O(n^2m)$, where $m$ is number of infected nodes, but it is easier to code.

Code

class Solution:
    def minMalwareSpread(self, graph, initial):
        n = len(graph)
        d = defaultdict(list)
        
        def dfs(node, par):
            for neib in range(n):
                if graph[node][neib] == 0: continue
                if neib in vis: continue
                vis.add(neib)
                d[neib].append(par)
                dfs(neib, par)
                
        for init in initial:
            vis = set(initial)
            dfs(init, init)
      
        res = [0] * n
        for key, val in d.items():
            if len(val) == 1: res[val[0]] += 1
        if max(res) == 0: return min(initial)
        
        return res.index(max(res))

Solution 3

We can use the idea of articulation points inside. The problem itself is not just to find these points, but also we need to keep sizes of components.

Complexity

Complexity of this code whould have been $O(n + E)$, but because data is given as adjacency matrix, it is $O(n^2)$.

Code

class Solution:
    def minMalwareSpread(self, graph, initial):
        def dfs(node, par):
            tin[node] = fup[node] = self.time
            self.time += 1
            flag = node in infected
            size = 1
            for child in range(n):
                if graph[node][child] == 0 or child == par: continue

                if tin[child] == 0:
                    s = dfs(child, node)
                    if s == 0:
                        flag = True
                    else:
                        size += s
                    
                    if fup[child] >= tin[node]: count[node] += s  #it means we are in articulation point.
                    fup[node] = min(fup[node], fup[child])
                else:
                    fup[node] = min(fup[node], tin[child])
                    
            return 0 if flag else size
        
        n, self.time = len(graph), 1
        tin, fup, count = [0]*n, [0]*n, [0]*n
        infected = set(initial)

        for u in initial:
            if fup[u] == 0: dfs(u, -1)
                
        return -max((count[u], -u) for u in initial)[1]