928. Minimize Malware Spread II

class Solution {
    // 1-2-3
    // 1, 2 infected
    // 1 3
    public int minMalwareSpread(int[][] graph, int[] initial) {
        int n = graph.length;
        Arrays.sort(initial);
        Set<Integer> mal = new HashSet<>();
        for (int i : initial) {
            mal.add(i);
        }
        
        int num = -1;
        int res = -1;
        for (int i : initial) {
            int save = 0;
            Set<Integer> visited = new HashSet<>();
            visited.add(i);
            
            for (int j = 0; j < n; j++) {
                if (i != j && graph[i][j] == 1) {
                    int tmp = dfs(j, mal, visited, graph);
                    if (tmp < 0) continue;
                    save += tmp;
                }
            }
            
            if (save > num) {
                num = save;
                res = i;
            }
        }
        
        return res;
    }
    
    private int dfs(int node, Set<Integer> mal, Set<Integer> visited, int[][] graph) {
        if (!visited.add(node)) return 0;
        if (mal.contains(node)) return -1;
        
        int save = 1;
        for (int j = 0; j < graph.length; j++) {
            if (node != j && graph[node][j] == 1) {
                int tmp = dfs(j, mal, visited, graph);
                if (tmp < 0) {
                    mal.add(node);
                    return -1;
                }
                
                save += tmp;
            }
        }
        
        return save;
    }
}