Update minimize-malware-spread-ii.cpp

Author: kamyu104

C++/minimize-malware-spread-ii.cpp

@@ -5,17 +5,17 @@ class Solution {
 public:
     int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {
         unordered_set<int> initial_set(initial.cbegin(), initial.cend());
-        UnionFind union_find(graph.size());
+        vector<int> clean;
         for (int i = 0; i < graph.size(); ++i) {
-            if (initial_set.count(i)) {
-                continue;
+            if (!initial_set.count(i)) {
+                clean.emplace_back(i);
             }
-            for (int j = i + 1; j < graph.size(); ++j) {
-                if (initial_set.count(j)) {
-                    continue;
-                }
-                if (graph[i][j] == 1) {
-                    union_find.union_set(i, j);
+        }
+        UnionFind union_find(graph.size());
+        for (int i = 0; i < clean.size(); ++i) {
+            for (int j = i + 1; j < clean.size(); ++j) {
+                if (graph[clean[i]][clean[j]] == 1) {
+                    union_find.union_set(clean[i], clean[j]);
                 }
             }
         }
@@ -26,26 +26,19 @@ class Solution {
 
         unordered_map<int, unordered_set<int>> shared_union;
         for (const auto& i: initial) {
-            for (int j = 0; j < graph.size(); ++j) {
-                if (initial_set.count(j)) {
-                    continue;
-                }
+            for (const auto& j : clean) {
                 if (graph[i][j] == 1) {
                     auto x = union_find.find_set(j);
                     shared_union[x].emplace(i);
                 }
             }
         }
-
         int result = numeric_limits<int>::max();
         int total = numeric_limits<int>::min();   
         for (const auto& i: initial) {
             unordered_set<int> lookup;
             int curr = 0;
-            for (int j = 0; j < graph.size(); ++j) {
-                if (initial_set.count(j)) {
-                    continue;
-                }
+            for (const auto& j : clean) {
                 if (graph[i][j] == 1) {
                     auto x = union_find.find_set(j);
                     if (shared_union[x].size() == 1 &&