Minimize Malware Spread LeetCode solution in Java

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

5
  public int minMalwareSpread(int[][] graph, int[] initial) {
6
    parent = new int[graph.length];
7
    size = new int[graph.length];
8
    for (int i = 0; i < parent.length; i++) {
9
      parent[i] = i;
10
      size[i] = 1;
11
    }
12

13
    for (int i = 0; i < graph.length; i++) {
14
      for (int j = 0; j < graph[0].length; j++) {
15
        if (i != j && graph[i][j] == 1) {
16
          unionHelper(i, j);
17
        }
18
      }
19
    }
20

21
    int[] imap = new int[graph.length];
22
    for (int val : initial) {
23
      int ri = find(val);
24
      imap[ri]++;
25
    }
26

27
    int ans_i = -1;
28
    int ans_size = 0;
29
    for (int i : initial) {
30
      int ri = find(i);
31
      if (imap[ri] == 1) {
32
        if (ans_i == -1) {
33
          ans_i = i;
34
          ans_size = size[ri];
35
        } else if (size[ri] > ans_size) {
36
          ans_i = i;
37
          ans_size = size[ri];
38
        } else if (size[ri] == ans_size) {
39
          if (i < ans_i) {
40
            ans_i = i;
41
            ans_size = size[ri];
42
          }
43
        }
44
      }
45
    }
46

47
    if (ans_i == -1) {
48
      ans_i = graph.length;
49
      for (int i : initial) {
50
        if (i < ans_i) {
51
          ans_i = i;
52
        }
53
      }
54
    }
55

56
    return ans_i;
57
  }
58

59
  int find(int x) {
60
    if (parent[x] == x) {
61
      return x;
62
    } else {
63
      parent[x] = find(parent[x]);
64
      return parent[x];
65
    }
66
  }
67

68
  void unionHelper(int x, int y) {
69
    int xl = find(x);
70
    int yl = find(y);
71

72
    if (size[xl] < size[yl]) {
73
      parent[xl] = yl;
74
      size[yl] += size[xl];
75
    } else {
76
      parent[yl] = xl;
77
      size[xl] += size[yl];
78
    }
79
  }
80
}