Minimize Malware Spread II LeetCode solution in Java

1
class Solution {
2
  int[] parent;
3
  int[] size;
4

5
  public int find(int x) {
6
    if (parent[x] == x) return x;
7
    int f = find(parent[x]);
8
    parent[x] = f;
9
    return f;
10
  }
11

12
  void merge(int x, int y) {
13
    if (size[x] > size[y]) {
14
      parent[y] = x;
15
      size[x] += size[y];
16
    } else {
17
      parent[x] = y;
18
      size[y] += size[x];
19
    }
20
  }
21

22
  public int minMalwareSpread(int[][] graph, int[] initial) {
23
    int n = graph.length;
24

25
    parent = new int[n];
26
    size = new int[n];
27
    // putting initially infected nodes in hashset to ignore them while making graph
28
    HashSet<Integer> hs = new HashSet<>();
29
    for (int a : initial) {
30
      hs.add(a);
31
    }
32
    // initializing Parent for DSU
33
    for (int i = 0; i < parent.length; i++) {
34
      parent[i] = i;
35
      size[i] = 1;
36
    }
37
    // constructing groups DSU
38
    for (int i = 0; i < graph.length; i++) {
39
      for (int j = 0; j < graph[0].length; j++) {
40
        if (graph[i][j] == 1 && !hs.contains(i) && !hs.contains(j)) {
41
          int p1 = find(i);
42
          int p2 = find(j);
43
          if (p1 != p2) {
44
            merge(p1, p2); // merging if they are not already
45
          }
46
        }
47
      }
48
    }
49
    // Storing initial Nodes vs parents of connected node...if there are multiple edge to same
50
    // component then we will get the same parent twice ...so we will take hashset for that purpose.
51
    HashMap<Integer, HashSet<Integer>> map = new HashMap<>();
52
    // We need to ensure increase the count whenever a parent is influenced by initial
53
    // Nodes...because if it influenced by more than one infectious node then it would still remain
54
    // infectious.
55
    int[] infected = new int[n];
56
    for (int e : initial) {
57
      map.put(e, new HashSet<>());
58
      for (int j = 0; j < n; j++) {
59
        // in adjaceny graph[i][i] is always one so we will ignore that
60
        if (!hs.contains(j) && e != j && graph[e][j] == 1) {
61
          int p = find(j);
62
          if (!map.get(e).contains(p)) {
63
            map.get(e).add(p);
64
            infected[p]++;
65
          }
66
        }
67
      }
68
    }
69
    int max = -1;
70
    int ans = -1;
71

72
    for (int e : initial) {
73
      HashSet<Integer> par = map.get(e);
74
      int total = 0;
75
      for (int p : par) {
76
        if (infected[p] == 1) { // add to total only if influenced by only one infectious node.
77
          total += size[p];
78
        }
79
      }
80
      if (total >= max) {
81
        if (max == total) {
82
          ans = Math.min(ans, e);
83
        } else {
84
          ans = e;
85
        }
86
        max = total;
87
      }
88
    }
89

90
    if (ans != -1) return ans;
91
    // for returining smallest element.
92
    Arrays.sort(initial);
93

94
    return initial[0];
95
  }
96
}