Solution.java

class Solution {
    private int[] p;
    private int[] size;

    public int minMalwareSpread(int[][] graph, int[] initial) {
        int n = graph.length;
        p = new int[n];
        size = new int[n];
        for (int i = 0; i < n; ++i) {
            p[i] = i;
            size[i] = 1;
        }
        boolean[] clean = new boolean[n];
        Arrays.fill(clean, true);
        for (int i : initial) {
            clean[i] = false;
        }
        for (int i = 0; i < n; ++i) {
            if (!clean[i]) {
                continue;
            }
            for (int j = i + 1; j < n; ++j) {
                if (clean[j] && graph[i][j] == 1) {
                    union(i, j);
                }
            }
        }
        int[] cnt = new int[n];
        Map<Integer, Set<Integer>> mp = new HashMap<>();
        for (int i : initial) {
            Set<Integer> s = new HashSet<>();
            for (int j = 0; j < n; ++j) {
                if (clean[j] && graph[i][j] == 1) {
                    s.add(find(j));
                }
            }
            for (int root : s) {
                cnt[root] += 1;
            }
            mp.put(i, s);
        }
        int mx = -1;
        int ans = 0;
        for (Map.Entry<Integer, Set<Integer>> entry : mp.entrySet()) {
            int i = entry.getKey();
            int t = 0;
            for (int root : entry.getValue()) {
                if (cnt[root] == 1) {
                    t += size[root];
                }
            }
            if (mx < t || (mx == t && i < ans)) {
                mx = t;
                ans = i;
            }
        }
        return ans;
    }

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

    private void union(int a, int b) {
        int pa = find(a);
        int pb = find(b);
        if (pa != pb) {
            size[pb] += size[pa];
            p[pa] = pb;
        }
    }
}