Minimize Malware Spread II.java

class Solution {
    int[] parent;
    int[] size;
   
    public int find(int x){
        if(parent[x]==x) return x;
        int f = find(parent[x]);
        parent[x] = f;
        return f;
    }
    
    void merge(int x, int y){
        if(size[x]>size[y]){
            parent[y] = x;
            size[x] += size[y];
        }else{
            parent[x] =y;
            size[y] +=size[x];
        }
    }
    
    public int minMalwareSpread(int[][] graph, int[] initial) {
        int n =graph.length;
        
        parent = new int[n];
        size= new int[n];
        // putting initially infected nodes in hashset to ignore them while making graph
        HashSet<Integer> hs = new HashSet<>();
        for(int a : initial){
            hs.add(a);
        }
        //initializing Parent for DSU
        for(int i=0;i<parent.length;i++){
            parent[i] = i;
            size[i]=1;
        }
        //constructing groups DSU
        for(int i=0;i<graph.length;i++){
            for(int j=0;j<graph[0].length;j++){
                if(graph[i][j]==1 && !hs.contains(i) && !hs.contains(j)){
                    int p1 = find(i);
                    int p2 = find(j);
                    if(p1!=p2){
                        merge(p1,p2); //merging if they are not already
                    }
                }
            }
        }
        //Storing initial Nodes vs parents of connected node...if there are multiple edge to same component then we will get the same parent twice ...so we will take hashset for that purpose.
        HashMap<Integer,HashSet<Integer>> map = new HashMap<>();
        //We need to ensure increase the count whenever a parent is influenced by initial Nodes...because if it influenced by more than one infectious node then it would still remain infectious.
        int[] infected = new int[n]; 
        for(int e: initial){
            map.put(e,new HashSet<>());
            for(int j=0;j<n;j++){
                // in adjaceny graph[i][i] is always one so we will ignore that
                if(!hs.contains(j) && e!=j && graph[e][j]==1){ 
                    int p = find(j);
                    if(!map.get(e).contains(p)){
                        map.get(e).add(p);
                        infected[p]++;
                    }
                }
            }
        }
        int max = -1;
        int ans =-1;
        
        for(int e:initial){
            HashSet<Integer> par = map.get(e);
            int total =0;
            for(int p: par){
                if(infected[p]==1){ //add to total only if influenced by only one infectious node.
                    total+=size[p];
                }
            }
            if(total>=max){
                if(max==total){
                    ans = Math.min(ans,e);
                }else{
                    ans =e;
                }
                max =total;
            }
        }
        
        if(ans!=-1) return ans;
        //for returining smallest element.
        Arrays.sort(initial);
        
        return initial[0];
        
    }
}