Minimize Malware Spread
Understanding the Problem
Given an undirected graph and a set of initially infected nodes, we need to find which single node, removed from the initial set, leaves the fewest total infected nodes after the malware finishes spreading.
Malware spreads across entire connected components. If a connected component contains at least one initially infected node, every node in that component eventually becomes infected. So the question reduces to: which initially infected node, when removed, saves the most nodes from infection?
Removing an initially infected node only saves its component if that node is the sole initially infected node in that component. If a component has two or more nodes from initial, removing one of them changes nothing, because the others still infect the entire component. The node to remove is therefore the one that is the only initially infected node in its component, and whose component is the largest.
Key Constraints:
2 <= n <= 300→ The adjacency matrix has at most 90,000 entries, so an O(n^2) scan over all edges is the dominant cost and runs comfortably within limits.graph[i][j] == graph[j][i]→ The graph is undirected, so Union Find can ignore edge direction.1 <= initial.length <= n→ The initial set can be the entire graph. Every node might be initially infected, so the per-component count must be handled even when no removal saves anything.
Approach 1: Brute Force (Simulate Each Removal)
Intuition
Try removing each node from initial one at a time, simulate the malware spread with BFS, and count how many nodes end up infected. Pick the removal that yields the fewest infections.
For each candidate node, build a modified initial set that excludes it, then run BFS from every remaining initially infected node to find all reachable nodes. The number of reachable nodes is the infection count for that removal.
Algorithm
- Sort the
initialarray so that ties are broken by smallest index: a later candidate replaces the current best only on a strictly smaller count. - For each node
removeNodeininitial: - Create a set of remaining infected nodes =
initialminusremoveNode. - Run BFS from all remaining infected nodes to count total infected nodes.
- Return the
removeNodethat resulted in the smallest infection count.
Example Walkthrough
Take a graph with two components: nodes 0 and 1 are connected, and nodes 2, 3, 4 are connected. The initial set is [0, 2], so each component starts with exactly one infected node. Removing 0 saves the 2-node component; removing 2 saves the 3-node component.
Code
- Time Complexity: O(k * n^2). For each of the k nodes in initial, we run a BFS that visits up to n nodes, checking all n neighbors for each. Since k can be up to n, worst case is O(n^3).
- Space Complexity: O(n). The infected array and BFS queue use O(n) space.
This runs a full BFS for every node in initial, yet the component structure is identical across all removals. Computing the connected components once removes the repeated work.
Approach 2: Connected Components (DFS)
Intuition
Malware spreads across entire connected components. If any node in a component is initially infected, all nodes in that component become infected, so the problem can be reasoned about one component at a time.
The deciding factor is, for each connected component, how many nodes from initial it contains:
- If a component contains zero nodes from
initial, it stays clean regardless of what we remove. - If a component contains exactly one node from
initial, removing that node saves the entire component from infection. - If a component contains two or more nodes from
initial, removing any one of them doesn't help because the others will still infect the whole component.
So the optimal node to remove is the one that's the sole infected node in the largest component. If no node is uniquely responsible for infecting its component, every removal saves zero nodes, and we return the smallest index in initial.
The component count of initial nodes fully determines the outcome of a removal. Within a single component, every infected starting node reaches the same set of nodes, so the order of infection and the number of starting nodes do not matter for the final infected set, only whether the count is zero.
If a component holds two or more initial nodes, removing one leaves at least one, and that one still infects the whole component, so savings are zero. If a component holds exactly one initial node, removing it leaves the component with no starting infection, and it stays clean. This is why savings equal the component size only in the count-equals-one case.
Algorithm
- Find all connected components using DFS. Assign each node a component ID and compute each component's size.
- For each component, count how many nodes from
initialbelong to it. - For each node in
initial: if it's the only initially infected node in its component, its "savings" equals the component size. - Return the node with the highest savings. Break ties by returning the smallest index.
Example Walkthrough
Use the same graph: nodes 0 and 1 form one component, nodes 2, 3, 4 form another, and initial = [0, 2]. The array below is componentId, the component label assigned to each of the five nodes.
Code
- Time Complexity: O(n^2). Finding connected components with DFS takes O(n^2) since we examine every edge in the adjacency matrix. The subsequent counting and comparison steps are O(n) or O(k).
- Space Complexity: O(n). We store component IDs, sizes, and initial counts, each requiring O(n) space.
The DFS approach requires explicit stack management and a separate component-labeling pass. Union Find builds the same component structure incrementally as it scans the edges.
Approach 3: Union Find (Optimal)
Intuition
Union Find groups connected nodes directly. Scan the adjacency matrix and union every pair of connected nodes. After processing all edges, nodes in the same component share the same root, so component sizes and the per-component count of initial nodes can be read off the roots in one pass each.
The decision logic is the same as Approach 2: find a node in initial that is the sole infected node in its component, and pick the one in the largest component. The only change is how components are identified.
Algorithm
- Initialize a Union Find structure with n nodes.
- Scan the adjacency matrix. For each edge
graph[i][j] == 1wherei < j, union nodesiandj. - After all unions, for each node in
initial, find its root (component representative). - Count how many nodes from
initialshare each root. Also track each component's total size. - Iterate through
initial(sorted). For each node whose component has exactly 1 initial node, its savings equals the component size. - Return the node with the highest savings, breaking ties by smallest index.
Example Walkthrough
Use the same graph again: nodes 0 and 1 connected, nodes 2, 3, 4 connected, initial = [0, 2]. The array below is parent, the Union Find parent pointer for each of the five nodes.
Code
- Time Complexity: O(n^2 * alpha(n)). Scanning the adjacency matrix takes O(n^2). Each union/find operation costs O(alpha(n)) amortized, which is effectively constant. Total is essentially O(n^2).
- Space Complexity: O(n). The parent and rank arrays each use O(n) space. Component size and initial count arrays also use O(n).