924 Minimize Malware Spread.py

Author: rajeevranjancom

924 Minimize Malware Spread.py

@@ -0,0 +1,101 @@
+#!/usr/bin/python3
+"""
+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.  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.
+
+Note that if a node was removed from the initial list of infected nodes, it may
+still be infected later as a result of the malware spread.
+
+Example 1:
+Input: graph = [[1,1,0],[1,1,0],[0,0,1]], initial = [0,1]
+Output: 0
+
+Example 2:
+Input: graph = [[1,0,0],[0,1,0],[0,0,1]], initial = [0,2]
+Output: 0
+
+Example 3:
+Input: graph = [[1,1,1],[1,1,1],[1,1,1]], initial = [1,2]
+Output: 1
+
+Note:
+1 < graph.length = graph[0].length <= 300
+0 <= graph[i][j] == graph[j][i] <= 1
+graph[i][i] = 1
+1 <= initial.length < graph.length
+0 <= initial[i] < graph.length
+"""
+from typing import List
+from collections import defaultdict
+
+
+class DisjointSet:
+    def __init__(self):
+        self.pi = {}
+
+    def union(self, x, y):
+        pi_x = self.find(x)
+        pi_y = self.find(y)
+        self.pi[pi_x] = pi_y
+
+    def find(self, x):
+        if x not in self.pi:
+            self.pi[x] = x
+        if self.pi[x] != x:
+            self.pi[x] = self.find(self.pi[x])
+        return self.pi[x]
+
+
+class Solution:
+    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:
+        """
+        DisjointSet. But how to use DisjointSet?
+
+        Ensure each component, the element points to a common ancestor.
+        The ancestor uniquely identify the component
+
+        Each component has size. If there are only one malware in the component,
+        then the component can be sanitized.
+        """
+        ds = DisjointSet()
+        n = len(graph)  # nbr matrix
+        for i in range(n):
+            for j in range(n):
+                if graph[i][j] == 1:
+                    ds.union(i, j)
+
+        counts = defaultdict(int)  # count of element in the component
+        for i in range(n):
+            counts[ds.find(i)] += 1
+
+        malware_counts = defaultdict(int)
+        for i in initial:
+            malware_counts[ds.find(i)] += 1
+
+        max_i = min(initial)
+        for i in initial:
+            pi = ds.find(i)
+            if malware_counts[pi] == 1:
+                max_count = counts[ds.find(max_i)]
+                if max_count < counts[pi]:
+                    max_i = i
+                elif max_count == counts[pi] and max_i > i:
+                    max_i = i
+
+        return max_i
+
+
+if __name__ == "__main__":
+    assert Solution().minMalwareSpread([[1,1,0],[1,1,0],[0,0,1]], [0,1]) == 0