Solution.java

class Solution {
    private int[] p;

    public int minMalwareSpread(int[][] graph, int[] initial) {
        int n = graph.length;
        p = new int[n];
        for (int i = 0; i < n; ++i) {
            p[i] = i;
        }
        int[] size = new int[n];
        Arrays.fill(size, 1);
        for (int i = 0; i < n; ++i) {
            for (int j = i + 1; j < n; ++j) {
                if (graph[i][j] == 1) {
                    int pa = find(i), pb = find(j);
                    if (pa == pb) {
                        continue;
                    }
                    p[pa] = pb;
                    size[pb] += size[pa];
                }
            }
        }
        int mi = Integer.MAX_VALUE;
        int res = initial[0];
        Arrays.sort(initial);
        for (int i = 0; i < initial.length; ++i) {
            int t = 0;
            Set<Integer> s = new HashSet<>();
            for (int j = 0; j < initial.length; ++j) {
                if (i == j) {
                    continue;
                }
                if (s.contains(find(initial[j]))) {
                    continue;
                }
                s.add(find(initial[j]));
                t += size[find(initial[j])];
            }
            if (mi > t) {
                mi = t;
                res = initial[i];
            }
        }
        return res;
    }

    private int find(int x) {
        if (p[x] != x) {
            p[x] = find(p[x]);
        }
        return p[x];
    }
}