Improved task 924.

Author: javadev

src/main/java/g0901_1000/s0924_minimize_malware_spread/Solution.java

@@ -1,68 +1,75 @@
 package g0901_1000.s0924_minimize_malware_spread;
 
 // #Hard #Array #Depth_First_Search #Breadth_First_Search #Matrix #Union_Find
-// #2022_03_29_Time_127_ms_(5.44%)_Space_129.2_MB_(17.35%)
-
-import java.util.HashMap;
-import java.util.HashSet;
+// #2022_04_29_Time_10_ms_(47.21%)_Space_55.1_MB_(83.27%)
 
 public class Solution {
-    private int[][] par;
+    private int[] size;
+    private int[] par;
 
     public int minMalwareSpread(int[][] graph, int[] initial) {
-        int n = graph.length;
-        par = new int[n][2];
-        for (int i = 0; i < par.length; i++) {
-            par[i][0] = i;
-            par[i][1] = 1;
+        size = new int[graph.length];
+        par = new int[graph.length];
+        for (int i = 0; i < graph.length; i++) {
+            size[i] = 1;
+            par[i] = i;
         }
-        for (int i = 0; i < n; i++) {
-            for (int j = i + 1; j < n; j++) {
+        // create groups
+        for (int i = 0; i < graph.length; i++) {
+            for (int j = 0; j < graph[0].length; j++) {
                 if (graph[i][j] == 1) {
-                    int li = find(i);
-                    int lj = find(j);
-                    if (li != lj) {
-                        par[li][0] = lj;
-                        par[lj][1] += par[li][1];
-                    }
+                    int p1 = find(i);
+                    int p2 = find(j);
+                    merge(p1, p2);
                 }
             }
         }
-        HashMap<Integer, HashSet<Integer>> map = new HashMap<>();
-        for (int val : initial) {
-            int lv = find(val);
-            if (!map.containsKey(lv)) {
-                map.put(lv, new HashSet<>());
-            }
-            map.get(lv).add(val);
+        // no of infected in group
+        int[] infected = new int[graph.length];
+        for (int e : initial) {
+            int p = find(e);
+            infected[p]++;
         }
+        int currSize = -1;
         int ans = -1;
-        int max = Integer.MIN_VALUE;
-        for (int val : initial) {
-            int lv = find(val);
-            if (map.get(lv).size() == 1) {
-                if (par[lv][1] > max) {
-                    max = par[lv][1];
-                    ans = val;
-                } else if (par[lv][1] == max && val < ans) {
-                    ans = val;
-                }
-            } else {
-                if (max < 0) {
-                    max = 0;
-                    ans = val;
-                } else if (max == 0 && val < ans) {
-                    ans = val;
+        for (int e : initial) {
+            int p = find(e);
+            if (infected[p] == 1 && size[p] >= currSize) {
+                if (size[p] > currSize) {
+                    ans = e;
+                } else {
+                    ans = Math.min(ans, e);
                 }
+                currSize = size[p];
+            }
+        }
+        // all groups have more than 1 infected node then return min value from initial
+        if (ans == -1) {
+            ans = initial[0];
+            for (int i = 0; i < initial.length; i++) {
+                ans = Math.min(ans, initial[i]);
             }
         }
         return ans;
     }
 
-    private int find(int x) {
-        if (par[x][0] == x) {
-            return x;
+    private void merge(int p1, int p2) {
+        if (p1 != p2) {
+            if (size[p1] > size[p2]) {
+                par[p2] = p1;
+                size[p1] += size[p2];
+            } else {
+                par[p1] = p2;
+                size[p2] += size[p1];
+            }
+        }
+    }
+
+    private int find(int u) {
+        if (par[u] == u) {
+            return u;
         }
-        return find(par[x][0]);
+        par[u] = find(par[u]);
+        return par[u];
     }
 }