Leetcode 928. Minimize Malware Spread II

Problem Explanation

In the given problem, we are provided with a network of nodes. Each node i is directly connected to another node j if and only if graph[i][j] = 1. Some nodes, represented by the array initial, are initially infected by malware.  Whenever two nodes are directly connected and at least one of those nodes is infected by malware, both nodes will be infected by malware.  This spread of malware will continue until no more nodes can be infected.

The problem is to remove one node from the initial list which would minimize the total number of nodes infected by malware after the spread stops. If multiple such nodes exist, we want to remove the one with the smallest index.

This problem can be solved by using the approach of Breadth-First Search (BFS).


We start with each node in the initial list and perform a breadth-first search, marking nodes as they are infected. We count the number of infected nodes. If the number of infected nodes is smaller than the smallest count seen so far, we update our chosen node to remove. We continue this way until we have tried all nodes in the initial list.

To see this in action, let's consider an example:

We are given: graph = [[1,1,0],[1,1,1],[0,1,1]], initial = [0,1]

If we remove node 0 from initial, the final total count of infected nodes will be 3, because all nodes are directly or indirectly connected to node 1 which is infected. However, if we remove node 1, the final total count of infected nodes will be 2, because only node 0 and node 1 are infected and no other node is connected to node 0 directly or indirectly.

So the final output will be 1, because the removal of this node causes the smallest final count of infected nodes.

Approach of the solution

Among all infected nodes in the initial list, the main idea is to find the node which, if removed, leads to the smallest total number of infected nodes. We use a breadth-first search from each infected node to find out the total number of infected nodes after the spread of malware stops. We repeat this process by excluding each of the infected nodes one by one from the BFS, thus effectively removing that node. Among all nodes, we return the one which resulted in the smallest number of infected nodes. If there are multiple such nodes, we choose the one with the smallest index.

Python Solution

3from collections import deque
4from typing import List
6class Solution:
7    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:        
8        def bfs(graph, removed, initial):
9            count = 0
10            visited = [False for _ in range(len(graph))]
11            visited[removed] = True
12            queue = deque(i for i in initial if i != removed)
14            while queue:
15                node = queue.popleft()
16                count += 1
17                for neighbour in range(len(graph)):
18                    if not visited[neighbour] and graph[node][neighbour] == 1:
19                        visited[neighbour] = True
20                        queue.append(neighbour)
21            return count
23        initial.sort()
24        infected_nodes = len(graph)
25        ans = -1  
27        for i in initial:
28            nodes = bfs(graph, i, initial)
29            if nodes < infected_nodes:
30                infected_nodes = nodes
31                ans = i
32        return ans

In the Python solution, we first define bfs() function which performs breadth-first search from all infected nodes excluding the one that is removed. This function returns the total number of infected nodes after the malware spread stopped. In the main function, we call this BFS for each of the initially infected nodes and return the node which gives us the smallest number infected nodes. If there are multiple such nodes, we return the one with the smallest index.## Java Solution

In Java, we use ArrayDeque for our queue and an array of booleans to keep track of visited nodes.

3import java.util.*;
5class Solution {
6    public int minMalwareSpread(int[][] graph, int[] initial) {
7        Arrays.sort(initial);        
8        int n = graph.length;
9        int ans = -1;
10        int infectedNodes = n;
12        for (int removed : initial) {
13            boolean[] visited = new boolean[n];
14            visited[removed] = true;
15            Deque<Integer> queue = new ArrayDeque<>();
16            for (int node : initial) {
17                if (node != removed) {
18                    queue.offer(node);
19                    visited[node] = true;
20                }
21            }
22            int count = 0;
23            while(!queue.isEmpty()) {
24                int node = queue.poll();
25                count++;
26                for (int nei = 0; nei < n; nei++) {
27                    if (graph[node][nei] == 1 && !visited[nei]) {
28                        visited[nei] = true;
29                        queue.offer(nei);
30                    }
31                }
32            }
33            if (count < infectedNodes) {       
34                infectedNodes = count;
35                ans = removed;
36            }
37        }
38        return ans;
39    }

In this Java solution, we iterate through each node in the initial array and set it as the removed node. For each node, we create a fresh visited array and queue, adding all nodes except the removed one. We then perform BFS from each of these nodes, updating the count of infected nodes. If this count is smaller than what we have seen so far, we update our answer to be the removed node.

Similar to the Python solution, the time complexity is O(n^3), as we are performing a full breadth-first search for each node in the graph, where n is the number of nodes.

JavaScript Solution

In JavaScript, we use an array for our queue and tracking visited nodes.

4 * @param {number[][]} graph
5 * @param {number[]} initial
6 * @return {number}
7 */
8var minMalwareSpread = function(graph, initial) {
9    initial.sort();
10    let numNodes = graph.length;
11    let ans = -1;
12    let infectedNodes = numNodes;
14    for (let removed of initial) {
15        let visited = Array(numNodes).fill(false);
16        visited[removed] = true;
17        let queue = [];
18        for (let node of initial) {
19            if (node !== removed) {
20                queue.push(node);
21                visited[node] = true;
22             }
23        }
24        let count = 0;
25        while(queue.length > 0) {
26            let node = queue.shift();
27            count++;
28            for (let next = 0; next < numNodes; next++) {
29                if (graph[node][next] === 1 && !visited[next]) {
30                    visited[next] = true;
31                    queue.push(next);
32                }
33            }
34        }
35        if (count < infectedNodes) {  
36            infectedNodes = count;
37            ans = removed;
38        }
39    }
40    return ans;

This JavaScript solution is structurally identical to the Java solution and the complexity remains the same. The only difference is syntactic, due to the syntactic differences between JavaScript and Java. The use of the Array function in JavaScript provides a convenient way of initializing a new array to the correct length and filling it with a default value.

Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.

TA 👨‍🏫