Problem Description

You have a network of n nodes represented as an n x n adjacency matrix called [graph](/problems/graph_intro). In this matrix, graph[i][j] == 1 means node i is directly connected to node j.

Initially, some nodes specified in the array initial are infected by malware. The malware spreads according to these rules:

• When two nodes are directly connected and at least one is infected, both become infected
• The spread continues until no more nodes can be infected

Let M(initial) denote the total number of nodes infected after the malware stops spreading.

Your task is to remove exactly one node from the initial array (removing the node and all its connections from the network) to minimize the final infection count M(initial).

Return the node that should be removed to minimize the infection. If multiple nodes would result in the same minimum infection count, return the node with the smallest index.

Example:

If you have a network where nodes are connected and initial = [0, 1] contains the initially infected nodes, you need to determine whether removing node 0 or node 1 would result in fewer total infections. The malware will spread through all connected components that contain infected nodes.

Key Points:

• You must remove exactly one initially infected node
• The removed node and all its connections are completely eliminated from the network
• After removal, malware spreads from the remaining initially infected nodes
• Choose the node whose removal minimizes total infection
• Break ties by choosing the smaller node index

Flowchart Walkthrough

First, let's pin down the algorithm using the Flowchart. Here's a step-by-step walkthrough:

Is it a graph?

• Yes: The problem explicitly states we have a network of n nodes represented as an n x n adjacency matrix graph, where graph[i][j] == 1 indicates a connection between nodes i and j.

Is it a tree?

• No: The graph is not necessarily a tree. It can have cycles and multiple connected components. The adjacency matrix represents a general graph structure.

Is the problem related to directed acyclic graphs?

• No: The adjacency matrix represents an undirected graph (malware spreads bidirectionally between connected nodes), and it may contain cycles.

Is the problem related to shortest paths?

• No: We're not finding shortest paths. Instead, we're analyzing how malware spreads through connected components and determining which node's removal minimizes infection spread.

Does the problem involve connectivity?

• Yes: The core of the problem is about connectivity. We need to understand which nodes are connected to determine how malware spreads through connected components. We also need to identify which connected components would be affected by removing each initially infected node.

Disjoint Set Union

• Yes: The flowchart suggests DSU (Disjoint Set Union) for connectivity problems, which aligns perfectly with the solution approach.

Conclusion: The flowchart leads us to use Disjoint Set Union (DSU/Union-Find) to handle the connectivity aspect of the problem. While DFS could also explore connected components, DSU is more efficient for:

1. Merging non-infected nodes into connected components
2. Quickly determining component sizes
3. Identifying which components are connected to infected nodes

The solution uses DSU to group non-infected nodes into connected components, then analyzes which initially infected node, when removed, would prevent the most infections by disconnecting these components from the malware spread.

Intuition

The key insight is to think about what happens when we remove an initially infected node from the network. We want to find which removal saves the most nodes from infection.

Let's break down what happens during malware spread:

1. Initially infected nodes will spread malware to all nodes in their connected components
2. If multiple infected nodes are in the same component, that component gets infected regardless of which one we remove
3. The only nodes we can "save" are those in components connected to exactly one initially infected node

This leads us to a crucial observation: we can only prevent infection in a connected component if it's threatened by exactly one initially infected node. If a component is connected to multiple infected nodes, removing any single infected node won't help - the other infected nodes will still infect that component.

So our strategy becomes:

1. First, identify all connected components that contain only non-infected nodes
2. For each initially infected node, find which non-infected components it connects to
3. Count how many initially infected nodes threaten each component
4. The benefit of removing an infected node is the total size of components that are only threatened by that node

Why use Union-Find? We need to:

• Group all non-infected nodes into their connected components efficiently
• Track the size of each component
• Quickly identify which component a non-infected node belongs to

The algorithm naturally follows:

1. Use Union-Find to merge all non-infected nodes that are connected to each other
2. For each initially infected node i, find all non-infected components it can reach
3. Count how many different infected nodes can reach each component
4. Calculate the benefit of removing each infected node (sum of sizes of components only it can reach)
5. Choose the infected node with maximum benefit (or smallest index if tied)

This approach elegantly handles the complex interaction between multiple infection sources and overlapping spread patterns.

Learn more about Depth-First Search, Breadth-First Search, Union Find and Graph patterns.

Solution Approach

The implementation follows our intuition by using Union-Find to manage connected components and carefully tracking which infected nodes threaten each component.

Step 1: Build Union-Find for Non-Infected Nodes

First, we create a Union-Find data structure and merge all non-infected nodes that are connected:

s = set(initial)  # Set of initially infected nodes
uf = UnionFind(n)
for i in range(n):
    if i not in s:  # Only process non-infected nodes
        for j in range(i + 1, n):
            if [graph](/problems/graph_intro)[i][j] and j not in s:  # Both nodes are non-infected and connected
                uf.union(i, j)

After this step, all non-infected nodes are grouped into their respective connected components.

Step 2: Map Infected Nodes to Threatened Components

For each initially infected node, we identify which non-infected components it can reach:

g = defaultdict(set)  # g[i] = set of component roots that infected node i can reach
cnt = Counter()       # cnt[root] = number of infected nodes that can reach this component

for i in initial:
    for j in range(n):
        if j not in s and [graph](/problems/graph_intro)[i][j]:  # j is non-infected and connected to i
            g[i].add(uf.find(j))  # Add j's component root to i's reachable set
  
    for root in g[i]:
        cnt[root] += 1  # Count how many infected nodes threaten this component

The g dictionary tells us which components each infected node threatens, while cnt tells us how many different infected nodes threaten each component.

Step 3: Calculate Benefit of Removing Each Infected Node

We evaluate each infected node to see how many nodes we can save by removing it:

ans, mx = 0, -1
for i in initial:
    t = sum(uf.get_size(root) for root in g[i] if cnt[root] == 1)
    if t > mx or (t == mx and i < ans):
        ans, mx = i, t

For each infected node i:

• We sum the sizes of all components that only i threatens (cnt[root] == 1)
• This gives us the total number of nodes we can save by removing i
• We keep track of the best choice, breaking ties by choosing the smaller index

Union-Find Implementation Details

The Union-Find class uses two optimizations:

1. Path Compression: In find(), we update self.p[x] to point directly to the root
2. Union by Size: When merging, we attach the smaller tree to the larger one

These optimizations ensure nearly constant time operations for union and find.

Time Complexity: O(n²) where n is the number of nodes, dominated by checking all pairs of nodes in the adjacency matrix.

Space Complexity: O(n) for the Union-Find structure and auxiliary data structures.

Example Walkthrough

Let's walk through a concrete example to illustrate the solution approach.

Consider a network with 5 nodes and the following connections:

graph = [[1,1,0,0,0],
         [1,1,1,0,0],
         [0,1,1,1,0],
         [0,0,1,1,1],
         [0,0,0,1,1]]
       
initial = [0, 2]  // Nodes 0 and 2 are initially infected

This forms a chain: 0 -- 1 -- 2 -- 3 -- 4

Step 1: Build Union-Find for Non-Infected Nodes

First, we identify non-infected nodes: {1, 3, 4}

Now we union non-infected nodes that are connected:

• Node 1 connects to node 2 (infected), so skip
• Node 3 connects to node 2 (infected) and node 4 (non-infected)
• Node 4 connects to node 3 (non-infected)
• Union(3, 4) creates component {3, 4} with size 2

After this step:

• Component containing node 1: {1} with size 1
• Component containing nodes 3,4: {3, 4} with size 2

Step 2: Map Infected Nodes to Threatened Components

For infected node 0:

• Node 0 connects to node 1 (non-infected)
• Component root of node 1 is 1
• g[0] = {1} (node 0 threatens component with root 1)

For infected node 2:

• Node 2 connects to node 1 (non-infected) and node 3 (non-infected)
• Component root of node 1 is 1
• Component root of node 3 is 3 (or 4, depending on union order)
• g[2] = {1, 3} (node 2 threatens components with roots 1 and 3)

Count how many infected nodes threaten each component:

• cnt[1] = 2 (threatened by both nodes 0 and 2)
• cnt[3] = 1 (threatened only by node 2)

Step 3: Calculate Benefit of Removing Each Infected Node

For removing node 0:

• Node 0 threatens component {1}
• But cnt[1] = 2, meaning component {1} is also threatened by node 2
• Even if we remove node 0, node 2 will still infect component {1}
• Benefit = 0 nodes saved

For removing node 2:

• Node 2 threatens components {1} and {3, 4}
• Component {1}: cnt[1] = 2, so it would still be infected by node 0
• Component {3, 4}: cnt[3] = 1, so only node 2 threatens it
• Benefit = size of component {3, 4} = 2 nodes saved

Result: Remove node 2 to save 2 nodes from infection.

Final infection outcomes:

• If we remove node 0: Nodes 1, 2, 3, 4 get infected (4 total)
• If we remove node 2: Only nodes 0, 1 get infected (2 total)

Therefore, the algorithm correctly returns node 2 as the optimal choice.

Solution Implementation

Time and Space Complexity

Time Complexity: O(n² × α(n))

The analysis breaks down as follows:

• Building the UnionFind structure for non-infected nodes: The nested loops iterate through all pairs (i, j) where i < j, performing union operations when both nodes are not in the initial infected set and are connected. This requires O(n²) iterations, with each union/find operation taking O(α(n)) amortized time, resulting in O(n² × α(n)).
• Processing initial infected nodes: For each node in initial (at most n nodes), we iterate through all n nodes to find connected components, performing find operations. This contributes O(n × |initial| × α(n)) which is bounded by O(n² × α(n)).
• Computing the answer: For each initial node, we sum up component sizes which takes O(|initial|) operations per node, bounded by O(n²).

The dominant term is O(n² × α(n)) where α(n) is the inverse Ackermann function, which grows extremely slowly and is effectively constant for all practical values of n.

Space Complexity: O(n)

The space usage includes:

• UnionFind data structure: O(n) for parent array p and size array size
• Set s for initial infected nodes: O(|initial|) which is bounded by O(n)
• Dictionary g storing connected components for each infected node: O(|initial| × n) in worst case, bounded by O(n²), but since we only store root representatives, it's effectively O(n) in practice
• Counter cnt: O(n) for counting occurrences

While the reference answer states O(n²) space complexity (likely accounting for the input graph), the additional space used by the algorithm itself is O(n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

Pitfall: Incorrectly Counting Components Threatened by Multiple Infected Nodes

The Problem: A common mistake is to assume that removing an infected node will always save all the non-infected nodes it connects to. This is incorrect when a non-infected component is threatened by multiple infected nodes.

Example Scenario: Consider a graph where:

• Component A (5 nodes) is connected to infected nodes 0 and 1
• Component B (3 nodes) is connected only to infected node 0

If we remove node 0:

• Component B is saved (3 nodes saved)
• Component A still gets infected from node 1 (0 nodes saved)
• Total saved: 3 nodes

The Mistake:

# WRONG: Counting all connected components without checking exclusivity
nodes_saved = sum(
    union_find.get_size(component_root) 
    for component_root in infected_to_components[infected_node]
)

This would incorrectly count Component A's 5 nodes as "saved" when removing node 0, even though node 1 would still infect it.

The Solution: Only count components that are threatened by exactly one infected node (the one we're considering removing):

# CORRECT: Only count components with exactly one threat
nodes_saved = sum(
    union_find.get_size(component_root) 
    for component_root in infected_to_components[infected_node] 
    if component_infection_count[component_root] == 1  # Key condition!
)

Pitfall: Building Union-Find with Infected Nodes Included

The Problem: Including infected nodes when building the Union-Find structure leads to incorrect component identification.

The Mistake:

# WRONG: Including infected nodes in Union-Find
for i in range(n):
    for j in range(i + 1, n):
        if graph[i][j]:  # Connects ALL nodes, including infected ones
            union_find.union(i, j)

This creates components that span across infected nodes, making it impossible to determine which non-infected components would be saved.

The Solution: Only union non-infected nodes to properly identify separate non-infected components:

# CORRECT: Only union non-infected nodes
for i in range(n):
    if i not in infected_set:  # Check i is not infected
        for j in range(i + 1, n):
            if graph[i][j] and j not in infected_set:  # Check j is not infected
                union_find.union(i, j)

Pitfall: Not Handling Edge Cases Properly

The Problem: Failing to handle cases where:

1. All initially infected nodes are isolated (not connected to any non-infected nodes)
2. The initial array contains only one node

The Mistake: Not initializing the answer properly or assuming there will always be nodes to save:

# WRONG: Uninitialized or incorrect default
ans = -1  # or ans = None
mx = 0

The Solution: Always initialize with a valid node from the initial array and handle the case where no nodes can be saved:

# CORRECT: Initialize with first node, handle -1 for "no savings"
best_node_to_remove = 0
max_nodes_saved = -1  # -1 ensures any non-negative savings will be better

# Or explicitly handle the edge case:
if not initial:
    return -1
best_node_to_remove = min(initial)  # Default to smallest index