Leetcode 928. Minimize Malware Spread II

(This problem is the same as Minimize Malware Spread, with the differences bolded.)

In a network of nodes, each node i is directly connected to another node j if and only if graph[i][j] = 1.

Some nodes initial are initially infected by malware.  Whenever two nodes are directly connected and at least one of those two nodes is infected by malware, both nodes will be infected by malware.  This spread of malware will continue until no more nodes can be infected in this manner.

Suppose M(initial) is the final number of nodes infected with malware in the entire network, after the spread of malware stops.

We will remove one node from the initial list, completely removing it and any connections from this node to any other node.  Return the node that if removed, would minimize M(initial).  If multiple nodes could be removed to minimize M(initial), return such a node with the smallest index.

 

Example 1:

Input: graph = [[1,1,0],[1,1,0],[0,0,1]], initial = [0,1]
Output: 0

Example 2:

Input: graph = [[1,1,0],[1,1,1],[0,1,1]], initial = [0,1]
Output: 1

Example 3:

Input: graph = [[1,1,0,0],[1,1,1,0],[0,1,1,1],[0,0,1,1]], initial = [0,1]
Output: 1

 

Note:

1. 1 < graph.length = graph[0].length <= 300
2. 0 <= graph[i][j] == graph[j][i] <= 1
3. graph[i][i] = 1
4. 1 <= initial.length < graph.length
5. 0 <= initial[i] < graph.length

Accepted

7,583

Submissions

18,933

-------------------------------------------------------------------------------------------------

核心思路是找到所有未感染区域有多少连通分量[g1,g2...,gn]，每个malware可能感染其中一个或者多个连通分量。如果某个连通分量被>=2个malware感染，隔离其中某一个malware并没有意义；但是如果某一个连通分量刚好被某一个malware感染，隔离这个malware有意义，这个连通分量里点的个数刚好贡献了这个malware的影响力。具体实现可以用并查集或者DFS，以下是并查集的实现：

class DSU:
    def find(self, x):
        return x if self.f[x] == x else self.find(self.f[x])

    def merge(self, x, y):
        rx = self.find(x)
        ry = self.find(y)
        self.f[ry] = rx

    def __init__(self, graph, initial):
        n = len(graph)
        self.f = [i for i in range(n)]
        for i in range(n):
            if (i not in initial):
                for j in range(i+1,n):
                    if (j not in initial and graph[i][j] == 1):
                        self.merge(i,j)
        self.cnt_dict = {} #每个不带毒连通区域有多少个节点
        for i in range(n):
            rt = self.find(i)
            self.cnt_dict[rt] = 1 if rt not in self.cnt_dict else self.cnt_dict[rt]+1

        self.malware_dict = {} #每个不带毒连通区域会被包含自己的多少个病毒感染
        for malware in initial:
            for j, adj in enumerate(graph[malware]):
                if (adj): #bug5: graph[i][j] == 1 and i != j
                    rt = self.find(j) #bug3: forget this line
                    self.malware_dict[rt] = 1 if rt not in self.malware_dict else self.malware_dict[rt]+1

class Solution(object):
    def minMalwareSpread(self, graph, initial):
        dsu = DSU(graph, set(initial))

        tuples = []
        for malware in initial:
            cnt = 0
            for i, adj in enumerate(graph[malware]): #bug2: for i in graph[malware]
                if (adj and malware != i): #bug4 graph[malware][i] == 1，自己感染自己不算，隔离了也没有用
                    rt = dsu.find(i)
                    if (rt in dsu.malware_dict and dsu.malware_dict[rt] == 1):
                        cnt += dsu.cnt_dict[rt]
            tuples.append((-cnt, malware))
        tuples.sort()
        return tuples[0][1]

s = Solution()