Leetcode 924. Minimize Malware Spread

Problem Explanation: This problem is about malware proliferation in a network of nodes, where each node is represented by a unique index. A two-dimensional array graph represents the network, where graph[i][j] = 1 implies that node i and j are directly connected. The malware spread is such that, if two nodes are directly connected, and at least one of them is infected, then the other also gets infected. The spread continues until no more nodes can be infected in this way. From the initial list of infected nodes, we are required to find a node whose removal would minimize the final number of nodes infected with malware in the entire network. If multiple nodes could be removed to achieve this, we need to return the one with the smallest index.

Approach to the Solution: The solution uses the Disjoint-Set (Union-Find) data structure to keep track of connected components and applies the concept of union by rank to reduce the time complexity. Each node represents a set, and in the beginning, every node is in its own set. Overtime, sets that are directly connect are merged together. The root of the set is the parent of every node in that set.

An example:

Consider the graph = [[1,1,0],[1,1,0],[0,0,1]] and initial = [0,1]. The network looks as follows:

1
2
3  0   2
4  |  
5  1

Initially, nodes 0 and 1 are infected. If we remove node 0, only node 1 remains infected. So, the final number of infected nodes is 1 (node 1).

The solution then maintains an array for tracking the parent of each node (id), another for the rank of each node (rank), and two additional arrays: ufSize to hold the size of the union each node is part of, and malwareCount to track the number of initial malware in each union. Finally, the solution iterates over the initial nodes and selects the one from the smallest union containing exactly 1 malware.

Python Solution:

1
2python
3class DSU:
4    def __init__(self, N):
5        self.p = list(range(N))
6
7    def find(self, x):
8        if self.p[x] != x:
9            self.p[x] = self.find(self.p[x])
10        return self.p[x]
11
12    def union(self, x, y):
13        xr = self.find(x)
14        yr = self.find(y)
15        self.p[xr] = yr
16
17
18class Solution:
19    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:
20        dsu = DSU(len(graph))
21        for i in range(len(graph)):
22            for j in range(i):
23                if graph[i][j]:
24                    dsu.union(i, j)
25        count = collections.Counter(dsu.find(u) for u in initial)
26        initial.sort()
27        ans = (float('inf'), float('inf'))
28        for node in initial:
29            root = dsu.find(node)
30            if count[root] == 1:
31                if ans == (float('inf'), float('inf')) or len([v for v in dsu.p if v == root]) < ans[0]:
32                    ans = (len([v for v in dsu.p if v == root]), node)
33        return ans[1] if ans != (float('inf'), float('inf')) else min(initial)

The above solution defines a DSU class to perform find and union operations and implements the logic to find the answer within the minMalwareSpread method of the Solution class.JavaScript Solution:

For the JavaScript solution, we will follow a similar approach. First, we need to create a class to handle union-find operations. Next, we'll initialize our disjoint-set and connect the nodes based on the graph. And eventually, we'll find the best node to remove from the initial array.

Here's the code:

1
2javascript
3class DSU {
4    constructor(N) {
5        this.parents = Array.from({length: N}, (_, i) => i);
6    }
7    
8    find(x) {
9        if (this.parents[x] !== x) {
10            this.parents[x] = this.find(this.parents[x]);
11        }
12        return this.parents[x];
13    }
14    
15    union(x, y) {
16        const xr = this.find(x);
17        const yr = this.find(y);
18        this.parents[xr] = yr;
19    }
20}
21
22var minMalwareSpread = function(graph, initial) {
23    const dsu = new DSU(graph.length);
24    for (let i = 0; i < graph.length; i++) {
25        for (let j = 0; j < i; j++) {
26            if (graph[i][j]) {
27                dsu.union(i, j);
28            }
29        }
30    }
31
32    const count = initial.reduce((res, u) => (res[dsu.find(u)] = (res[dsu.find(u)] || 0) + 1, res), {});
33    initial.sort();
34    let ans = [Infinity, Infinity];
35    for (let node of initial) {
36        const root = dsu.find(node);
37        if (count[root] === 1) {
38            const cur = [dsu.parents.filter(v => v === root).length, node];
39            if (cur[0] < ans[0] || (cur[0] === ans[0] && cur[1] < ans[1])) {
40                ans = cur;
41            }
42        }
43    }
44    return ans[1] < Infinity ? ans[1] : Math.min(...initial);
45};

Java Solution:

In Java, our approach remains similar. We'll create a DSU class with methods for find and union operations. In the minMalwareSpread function, we'll use this class to determine the best node to remove.

The java solution is:

1
2java
3class DSU {
4    int[] parent;
5    public DSU(int N) {
6        parent = new int[N];
7        for (int i = 0; i < N; ++i)
8            parent[i] = i;
9    }
10    public int find(int x) {
11        if (parent[x] != x) parent[x] = find(parent[x]);
12        return parent[x];
13    }
14    public void union(int x, int y) {
15        parent[find(x)] = find(y);
16    }
17}
18
19class Solution {
20    public int minMalwareSpread(int[][] graph, int[] initial) {
21        int N = graph.length;
22        DSU dsu = new DSU(N);
23        for (int i = 0; i < N; ++i)
24            for (int j = 0; j < i; ++j)
25                if (graph[i][j] == 1)
26                    dsu.union(i, j);
27
28        int[] count = new int[N];
29        for (int node: initial)
30            count[dsu.find(node)]++;
31
32        Arrays.sort(initial);
33        int ans = initial[0], ansSize = -1;
34        for (int node: initial) {
35            int root = dsu.find(node);
36            if (count[root] == 1) {
37                int rootSize = 0;
38                for (int p: dsu.parent) if (dsu.find(p) == root) rootSize++;
39                if (rootSize > ansSize) {
40                    ansSize = rootSize;
41                    ans = node;
42                }
43            }
44        }
45        return ans;
46    }
47}

In the minMalwareSpread function, we can define the DSU, perform the union, and then find the smallest union containing exactly 1 malware. This approach extracts the optimal node for removal to minimize malware spread.


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 👨‍🏫