924. Minimize Malware Spread

1class Solution {
2    int[] parent, size;
3
4    public int minMalwareSpread(int[][] graph, int[] initial) {
5        int n = graph.length;
6        parent = new int[n];
7        size = new int[n];
8
9        // Initialize Union-Find
10        for (int i = 0; i < n; i++) {
11            parent[i] = i;
12            size[i] = 1;
13        }
14
15        // Build components from adjacency matrix
16        for (int i = 0; i < n; i++) {
17            for (int j = i + 1; j < n; j++) {
18                if (graph[i][j] == 1) {
19                    union(i, j);
20                }
21            }
22        }
23
24        // Count infected nodes per component
25        Map<Integer, List<Integer>> infected = new HashMap<>();
26        for (int node : initial) {
27            int root = find(node);
28            infected.computeIfAbsent(root, k -> new ArrayList<>()).add(node);
29        }
30
31        // Find best node to remove
32        Arrays.sort(initial);
33        int bestNode = initial[0];
34        int maxSaved = 0;
35
36        for (int node : initial) {
37            int root = find(node);
38            List<Integer> list = infected.get(root);
39
40            if (list.size() == 1) {
41                int saved = size[root];
42                if (saved > maxSaved || (saved == maxSaved && node < bestNode)) {
43                    maxSaved = saved;
44                    bestNode = node;
45                }
46            }
47        }
48
49        return bestNode;
50    }
51
52    int find(int x) {
53        if (parent[x] != x) parent[x] = find(parent[x]);
54        return parent[x];
55    }
56
57    void union(int x, int y) {
58        int px = find(x), py = find(y);
59        if (px != py) {
60            if (size[px] < size[py]) { int t = px; px = py; py = t; }
61            parent[py] = px;
62            size[px] += size[py];
63        }
64    }
65}