MinimizeMalwareSpread.java

/*https://leetcode.com/problems/minimize-malware-spread/*/

class Solution {
    int[] colors;
    public int minMalwareSpread(int[][] graph, int[] initial) {
        int N = graph.length;
        colors = new int[N];
        Arrays.fill(colors,-1);
        int color = 0;
        
        //coloring components
        for (int node = 0; node < N; ++node)
            if (colors[node] == -1)
                dfs(graph,node,color++);
        
        //finding size of each component
        int[] size = new int[color];
        for (int c : colors)
            ++size[c];
        
        //find unique colors
        int[] colorCount = new int[color];
        for (int node : initial)
            ++colorCount[colors[node]];
        
        int result = Integer.MAX_VALUE;
        for (int node : initial)
        {
            int c = colors[node];
            if (colorCount[c] == 1) //if the node is uniquely colored
            {
                if (result == Integer.MAX_VALUE) //if nothing is found yet
                    result = node; //store in result
                else if (size[c] > size[colors[result]]) //if size is larger, update result
                    result = node;
                else if (size[c] == size[colors[result]] && node < result) //if index is smaller, update result
                    result = node;
            }
        }
        
        if (result == Integer.MAX_VALUE) //if no such node could be found
        {
            Arrays.sort(initial);
            return initial[0]; //return the minimum node
        }
        
        return result;
    }
    private void dfs(int[][] graph, int src, int color)
    {
        colors[src] = color;
        for (int neighbour = 0; neighbour < graph.length; ++neighbour)
            if (graph[src][neighbour] == 1 && colors[neighbour] == -1)
                dfs(graph, neighbour, color);
    }
}