feat: add solutions to lc problem: No.0928

No.0928.Minimize Malware Spread II

Author: yanglbme

Diff for: solution/0900-0999/0928.Minimize Malware Spread II/README.md

@@ -55,7 +55,13 @@
 
 <!-- 这里可写通用的实现逻辑 -->
 
-并查集。
+逆向思维并查集。对于本题，先遍历所有未被感染的节点（即不在 initial 列表的节点），构造并查集，并且在集合根节点维护 size，表示当前集合的节点数。
+
+然后找到只被一个 initial 节点感染的集合，求得感染节点数的最小值。
+
+> 被某个 initial 节点感染的集合，节点数越多，若移除此 initial 节点，感染的节点数就越少。
+
+以下是并查集的几个常用模板。
 
 模板 1——朴素并查集：
 
@@ -114,12 +120,6 @@ p[find(a)] = find(b)
 d[find(a)] = distance
 ```
 
-对于本题，先遍历所有未被感染的节点（即不在 initial 列表的节点），构造并查集，并且在集合根节点维护 size，表示当前集合的节点数。
-
-然后找到只被一个 initial 节点感染的集合，求得感染节点数的最小值。
-
-> 被某个 initial 节点感染的集合，节点数越多，若移除此 initial 节点，感染的节点数就越少。
-
 <!-- tabs:start -->
 
 ### **Python3**
@@ -129,56 +129,43 @@ d[find(a)] = distance
 ```python
 class Solution:
     def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:
-        n = len(graph)
-        p = list(range(n))
-        size = [1] * n
-
         def find(x):
             if p[x] != x:
                 p[x] = find(p[x])
             return p[x]
 
+        def union(a, b):
+            pa, pb = find(a), find(b)
+            if pa != pb:
+                size[pb] += size[pa]
+                p[pa] = pb
+
+        n = len(graph)
+        p = list(range(n))
+        size = [1] * n
         clean = [True] * n
         for i in initial:
             clean[i] = False
-
         for i in range(n):
             if not clean[i]:
                 continue
             for j in range(i + 1, n):
-                if not clean[j]:
-                    continue
-                if graph[i][j] == 1:
-                    pa, pb = find(i), find(j)
-                    if pa == pb:
-                        continue
-                    p[pa] = pb
-                    size[pb] += size[pa]
-
+                if clean[j] and graph[i][j] == 1:
+                    union(i, j)
         cnt = Counter()
         mp = {}
         for i in initial:
-            s = set()
-            for j in range(n):
-                if not clean[j]:
-                    continue
-                if graph[i][j] == 1:
-                    s.add(find(j))
-            for e in s:
-                cnt[e] += 1
+            s = {find(j) for j in range(n) if clean[j] and graph[i][j] == 1}
+            for root in s:
+                cnt[root] += 1
             mp[i] = s
 
-        mx = -1
-        res = 0
+        mx, ans = -1, 0
         for i, s in mp.items():
-            t = 0
-            for e in s:
-                if cnt[e] == 1:
-                    t += size[e]
-            if mx < t or (mx == t and i < res):
-                mx = t
-                res = i
-        return res
+            t = sum(size[root] for root in s if cnt[root] == 1)
+            if mx < t or mx == t and i < ans:
+                mx, ans = t, i
+        return ans
 ```
 
 ### **Java**
@@ -188,11 +175,12 @@ class Solution:
 ```java
 class Solution {
     private int[] p;
+    private int[] size;
 
     public int minMalwareSpread(int[][] graph, int[] initial) {
         int n = graph.length;
         p = new int[n];
-        int[] size = new int[n];
+        size = new int[n];
         for (int i = 0; i < n; ++i) {
             p[i] = i;
             size[i] = 1;
@@ -207,16 +195,8 @@ class Solution {
                 continue;
             }
             for (int j = i + 1; j < n; ++j) {
-                if (!clean[j]) {
-                    continue;
-                }
-                if (graph[i][j] == 1) {
-                    int pa = find(i), pb = find(j);
-                    if (pa == pb) {
-                        continue;
-                    }
-                    p[pa] = pb;
-                    size[pb] += size[pa];
+                if (clean[j] && graph[i][j] == 1) {
+                    union(i, j);
                 }
             }
         }
@@ -225,34 +205,31 @@ class Solution {
         for (int i : initial) {
             Set<Integer> s = new HashSet<>();
             for (int j = 0; j < n; ++j) {
-                if (!clean[j]) {
-                    continue;
-                }
-                if (graph[i][j] == 1) {
+                if (clean[j] && graph[i][j] == 1) {
                     s.add(find(j));
                 }
             }
-            for (int e : s) {
-                cnt[e] += 1;
+            for (int root : s) {
+                cnt[root] += 1;
             }
             mp.put(i, s);
         }
         int mx = -1;
-        int res = 0;
+        int ans = 0;
         for (Map.Entry<Integer, Set<Integer>> entry : mp.entrySet()) {
             int i = entry.getKey();
             int t = 0;
-            for (int e : entry.getValue()) {
-                if (cnt[e] == 1) {
-                    t += size[e];
+            for (int root : entry.getValue()) {
+                if (cnt[root] == 1) {
+                    t += size[root];
                 }
             }
-            if (mx < t || (mx == t && i < res)) {
+            if (mx < t || (mx == t && i < ans)) {
                 mx = t;
-                res = i;
+                ans = i;
             }
         }
-        return res;
+        return ans;
     }
 
     private int find(int x) {
@@ -261,6 +238,15 @@ class Solution {
         }
         return p[x];
     }
+
+    private void union(int a, int b) {
+        int pa = find(a);
+        int pb = find(b);
+        if (pa != pb) {
+            size[pb] += size[pa];
+            p[pa] = pb;
+        }
+    }
 }
 ```
 
@@ -270,100 +256,101 @@ class Solution {
 class Solution {
 public:
     vector<int> p;
+    vector<int> size;
 
     int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {
         int n = graph.size();
-        vector<int> size(n, 1);
-        for (int i = 0; i < n; ++i) p.push_back(i);
+        p.resize(n);
+        size.resize(n);
+        for (int i = 0; i < n; ++i) p[i] = i;
+        fill(size.begin(), size.end(), 1);
         vector<bool> clean(n, true);
         for (int i : initial) clean[i] = false;
         for (int i = 0; i < n; ++i)
         {
             if (!clean[i]) continue;
             for (int j = i + 1; j < n; ++j)
-            {
-                if (!clean[j]) continue;
-                if (graph[i][j])
-                {
-                    int pa = find(i), pb = find(j);
-                    if (pa == pb) continue;
-                    p[pa] = pb;
-                    size[pb] += size[pa];
-                }
-            }
+                if (clean[j] && graph[i][j] == 1) merge(i, j);
         }
         vector<int> cnt(n, 0);
         unordered_map<int, unordered_set<int>> mp;
         for (int i : initial)
         {
             unordered_set<int> s;
             for (int j = 0; j < n; ++j)
-            {
-                if (!clean[j]) continue;
-                if (graph[i][j]) s.insert(find(j));
-            }
+                if (clean[j] && graph[i][j] == 1) s.insert(find(j));
             for (int e : s) ++cnt[e];
             mp[i] = s;
         }
-        int mx = -1;
-        int res = 0;
-        for (auto item : mp)
+        int mx = -1, ans = 0;
+        for (auto& [i, s] : mp)
         {
-            int i = item.first;
             int t = 0;
-            for (int e : item.second)
-            {
-                if (cnt[e] == 1) t += size[e];
-            }
-            if (mx < t || (mx == t && i < res))
+            for (int root : s)
+                if (cnt[root] == 1)
+                    t += size[root];
+            if (mx < t || (mx == t && i < ans))
             {
                 mx = t;
-                res = i;
+                ans = i;
             }
         }
-        return res;
+        return ans;
     }
 
     int find(int x) {
         if (p[x] != x) p[x] = find(p[x]);
         return p[x];
     }
+
+    void merge(int a, int b) {
+        int pa = find(a), pb = find(b);
+        if (pa != pb)
+        {
+            size[pb] += size[pa];
+            p[pa] = pb;
+        }
+    }
 };
 ```
 
 ### **Go**
 
 ```go
-var p []int
-
 func minMalwareSpread(graph [][]int, initial []int) int {
 	n := len(graph)
-	p = make([]int, n)
+	p := make([]int, n)
 	size := make([]int, n)
 	clean := make([]bool, n)
 	for i := 0; i < n; i++ {
-		p[i] = i
-		size[i] = 1
-		clean[i] = true
+		p[i], size[i], clean[i] = i, 1, true
 	}
 	for _, i := range initial {
 		clean[i] = false
 	}
+
+	var find func(x int) int
+	find = func(x int) int {
+		if p[x] != x {
+			p[x] = find(p[x])
+		}
+		return p[x]
+	}
+	union := func(a, b int) {
+		pa, pb := find(a), find(b)
+		if pa != pb {
+			size[pb] += size[pa]
+			p[pa] = pb
+		}
+	}
+
 	for i := 0; i < n; i++ {
 		if !clean[i] {
 			continue
 		}
 		for j := i + 1; j < n; j++ {
-			if !clean[j] {
-				continue
-			}
-			if graph[i][j] == 1 {
-				pa, pb := find(i), find(j)
-				if pa == pb {
-					continue
-				}
-				p[pa] = pb
-				size[pb] += size[pa]
+			if clean[j] && graph[i][j] == 1 {
+				union(i, j)
 			}
 		}
 	}
@@ -372,38 +359,28 @@ func minMalwareSpread(graph [][]int, initial []int) int {
 	for _, i := range initial {
 		s := make(map[int]bool)
 		for j := 0; j < n; j++ {
-			if !clean[j] {
-				continue
-			}
-			if graph[i][j] == 1 {
+			if clean[j] && graph[i][j] == 1 {
 				s[find(j)] = true
 			}
 		}
-		for e, _ := range s {
-			cnt[e]++
+		for root, _ := range s {
+			cnt[root]++
 		}
 		mp[i] = s
 	}
-	mx, res := -1, 0
+	mx, ans := -1, 0
 	for i, s := range mp {
 		t := 0
-		for e, _ := range s {
-			if cnt[e] == 1 {
-				t += size[e]
+		for root, _ := range s {
+			if cnt[root] == 1 {
+				t += size[root]
 			}
 		}
-		if mx < t || (mx == t && i < res) {
-			mx, res = t, i
+		if mx < t || (mx == t && i < ans) {
+			mx, ans = t, i
 		}
 	}
-	return res
-}
-
-func find(x int) int {
-	if p[x] != x {
-		p[x] = find(p[x])
-	}
-	return p[x]
+	return ans
 }
 ```