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 with malware. The malware spreads according to these rules:

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

The function M(initial) represents the total number of nodes infected after the malware finishes spreading.

Your task is to remove exactly one node from the initial array to minimize M(initial). You need to find which node to remove that results in the smallest number of total infections.

Key points to consider:

• You must remove exactly one initially infected node
• After removal, that node starts as uninfected (but could still get infected later if connected to other infected nodes)
• If multiple nodes result in the same minimum infection count, return the one with the smallest index
• The goal is to minimize the final total number of infected nodes in the network

For example, if removing node 2 from initial results in 5 total infections, and removing node 3 results in 4 total infections, you would choose to remove node 3. If both result in the same number of infections, you'd choose the one with the smaller 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 where graph[i][j] == 1 indicates nodes i and j are connected.

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 problem deals with an undirected graph (the adjacency matrix is symmetric) and doesn't involve directed edges or acyclic properties.

Is the problem related to shortest paths?

• No: We're not finding shortest paths between nodes. Instead, we're tracking malware spread and trying to minimize the total number of infected nodes.

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 need to identify connected components and count infected nodes within each component.

Disjoint Set Union

• Yes: DSU (Union-Find) is the suggested approach for handling connectivity problems efficiently.

Conclusion: The flowchart suggests using Disjoint Set Union (DSU) for this problem. While DFS could also solve connectivity problems, DSU is more efficient for:

1. Finding connected components in the graph
2. Determining the size of each connected component
3. Identifying which infected nodes belong to which component

The solution uses DSU to group nodes into connected components, then analyzes which initially infected node, when removed, would prevent the most infections by considering components with only one initially infected node.

Intuition

Let's think about how malware spreads in this network. Once a node is infected, it will infect all nodes in its connected component - there's no way to stop the spread within a connected component. This is the key insight.

Now, when we remove one initially infected node, what happens? We need to consider different scenarios based on how many initially infected nodes are in each connected component:

Case 1: A connected component has 0 initially infected nodes

• This component stays clean regardless of which node we remove
• No impact on our decision

Case 2: A connected component has exactly 1 initially infected node

• If we remove this node, the entire component stays clean!
• We save all nodes in this component from infection
• This is where we can make a difference

Case 3: A connected component has 2 or more initially infected nodes

• Even if we remove one infected node, others will still infect the entire component
• The whole component gets infected anyway
• Removing any of these nodes doesn't help

This leads us to our strategy: We should focus on Case 2 components - those with exactly one initially infected node. By removing that single infected node, we prevent the entire component from being infected.

To maximize our impact, we want to:

1. Find all connected components using Union-Find (efficient for connectivity)
2. Count how many initially infected nodes are in each component
3. Among components with exactly 1 infected node, choose to remove the infected node from the largest component (saves the most nodes)
4. If there's a tie in component size, pick the node with the smallest index

If no component has exactly 1 infected node (all have 2+), then removing any node won't reduce infections - we just return the smallest indexed node as required.

The Union-Find data structure is perfect here because it efficiently:

• Groups nodes into connected components
• Tracks the size of each component
• Allows us to quickly identify which component each infected node belongs to

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

Solution Approach

Let's implement our strategy using the Union-Find data structure to efficiently manage connected components.

Step 1: Build the Union-Find Structure

First, we create a Union-Find class with path compression and union by size optimization:

• find(x): Finds the root of node x with path compression
• union(a, b): Merges two components, always attaching the smaller component to the larger one
• get_size(root): Returns the size of the component with given root

class UnionFind:
    def __init__(self, n):
        self.p = list(range(n))  # parent array, initially each node is its own parent
        self.size = [1] * n      # size of each component

Step 2: Build Connected Components

We iterate through the adjacency matrix and union all connected nodes:

for i in range(n):
    for j in range(i + 1, n):
        if [graph](/problems/graph_intro)[i][j] == 1:
            uf.union(i, j)

After this step, all nodes in the same connected component will have the same root.

Step 3: Count Initially Infected Nodes per Component

We count how many initially infected nodes belong to each component:

cnt = Counter(uf.find(x) for x in initial)

Here, cnt[root] tells us how many initially infected nodes are in the component with root root.

Step 4: Find the Best Node to Remove

We examine each initially infected node and check if it's the only infected node in its component:

ans, mx = n, 0  # ans tracks the node to remove, mx tracks max nodes we can save

for x in initial:
    root = uf.find(x)
    if cnt[root] > 1:  # Multiple infected nodes in this component
        continue       # Removing x won't help
  
    # Only one infected node (x) in this component
    sz = uf.get_size(root)  # Size of this component
  
    if sz > mx or (sz == mx and x < ans):
        ans = x
        mx = sz

The logic here is:

• If cnt[root] > 1, there are other infected nodes in this component, so removing x won't prevent the component from being infected
• If cnt[root] == 1, removing x will save the entire component of size sz
• We update our answer if we find a larger component to save, or same size but smaller node index

Step 5: Handle Edge Case

If no component has exactly 1 infected node (meaning ans was never updated from n), we return the minimum indexed infected node:

return min(initial) if ans == n else ans

Time Complexity: O(n²) for building the graph connections + O(m × α(n)) for Union-Find operations where m is the number of initially infected nodes and α is the inverse Ackermann function (practically constant).

Space Complexity: O(n) for the Union-Find structure and counting array.

Example Walkthrough

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

Example Setup:

• Network with 6 nodes (0-5)
• Initially infected nodes: initial = [0, 1, 4]
• Graph connections:
  • Nodes 0, 1, 2 form a connected component (all connected to each other)
  • Nodes 3, 4 form another connected component
  • Node 5 is isolated

Adjacency Matrix:
    0 1 2 3 4 5
0 [ 0 1 1 0 0 0 ]
1 [ 1 0 1 0 0 0 ]
2 [ 1 1 0 0 0 0 ]
3 [ 0 0 0 0 1 0 ]
4 [ 0 0 0 1 0 0 ]
5 [ 0 0 0 0 0 0 ]

Step 1: Build Union-Find and identify components

After processing all edges:

• Component A: {0, 1, 2} with root 0, size = 3
• Component B: {3, 4} with root 3, size = 2
• Component C: {5} with root 5, size = 1

Step 2: Count infected nodes per component

• Component A has 2 infected nodes: 0 and 1
• Component B has 1 infected node: 4
• Component C has 0 infected nodes

Step 3: Evaluate each initially infected node

For node 0:

• Root = 0, component size = 3
• Component has 2 infected nodes (0 and 1)
• Since count > 1, removing node 0 won't help (node 1 will still infect the component)
• Skip this option

For node 1:

• Root = 0, component size = 3
• Component has 2 infected nodes (0 and 1)
• Since count > 1, removing node 1 won't help (node 0 will still infect the component)
• Skip this option

For node 4:

• Root = 3, component size = 2
• Component has exactly 1 infected node (only node 4)
• If we remove node 4, we save the entire component of size 2
• This is a valid candidate: saves 2 nodes

Step 4: Make the decision

• Only node 4 can actually reduce infections
• Removing node 4 saves 2 nodes (the entire Component B)
• Final answer: Remove node 4

Verification of outcomes:

• If we remove node 0: Nodes 1,2 get infected (Component A), node 4 infects node 3 (Component B) → Total: 4 infected
• If we remove node 1: Nodes 0,2 get infected (Component A), node 4 infects node 3 (Component B) → Total: 4 infected
• If we remove node 4: Nodes 0,1,2 get infected (Component A), nodes 3,4 stay clean → Total: 3 infected ✓

The algorithm correctly identifies that removing node 4 minimizes the total infection count.

Solution Implementation

Time and Space Complexity

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

The time complexity is dominated by the graph traversal and union-find operations:

• The nested loops iterate through all pairs of nodes: O(n²) iterations
• For each edge where graph[i][j] = 1, we perform a union operation
• Each union operation involves two find operations, which have amortized complexity of O(α(n)) due to path compression
• Finding roots for all initial nodes: O(|initial| × α(n)) where |initial| ≤ n
• The final loop through initial nodes: O(|initial|) operations

Therefore, the overall time complexity 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 complexity consists of:

• Union-Find data structure stores two arrays p and size, each of size n: O(n)
• Counter for storing root counts: O(|initial|) which is at most O(n)
• Recursive call stack for find operation with path compression: O(log n) in worst case before compression
• Other variables use constant space: O(1)

The total space complexity is O(n).

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

Common Pitfalls

Pitfall 1: Misunderstanding What "Removing" a Node Means

The Problem: Many developers incorrectly interpret "removing a node" as deleting it entirely from the graph, including all its edges. This leads to complicated graph reconstruction logic and incorrect results.

What Actually Happens: When we "remove" a node from the initial array, we're only changing its initial infection status from infected to clean. The node and its connections remain in the graph. It can still get infected later if connected to other infected nodes.

Incorrect Approach:

# WRONG: Trying to delete the node from the graph
def remove_node(graph, node):
    n = len(graph)
    for i in range(n):
        graph[node][i] = 0
        graph[i][node] = 0

Correct Understanding: The node stays in the graph with all its connections intact. We're just choosing not to start with it as infected.

Pitfall 2: Incorrectly Handling Components with Multiple Infected Nodes

The Problem: When a connected component has multiple initially infected nodes, some might think removing one of them reduces the infection count by the component size.

Why It's Wrong: If a component has nodes A, B, C, D where A and B are initially infected:

• Even if we remove A from initial, B is still infected
• B will spread to all nodes in the component including A
• The entire component still gets infected

Incorrect Logic:

# WRONG: Counting saved nodes even when multiple infected in component
for node in initial:
    component_root = uf.find(node)
    saved_nodes = uf.get_size(component_root)  # Wrong! Doesn't check if others are infected
    # ... update best choice

Correct Logic:

# RIGHT: Only count savings when exactly one infected node in component
if infected_count_per_component[component_root] > 1:
    continue  # Can't save this component

Pitfall 3: Not Handling the "All Components Have Multiple Infections" Case

The Problem: If every connected component containing initially infected nodes has 2+ infected nodes, removing any single node won't prevent any component from being infected. The code might return an invalid result or crash.

Example Scenario:

graph = [[1,1,0], [1,1,0], [0,0,1]]
initial = [0, 1]  # Both in same component
# Component {0,1} has 2 infected nodes
# Component {2} has 0 infected nodes
# No matter which we remove, component {0,1} stays infected

Solution: Always have a fallback to return the minimum indexed node when no component can be saved:

if best_node_to_remove == n:  # Never updated, no single-infected component found
    return min(initial)  # Return smallest index as tiebreaker

Pitfall 4: Incorrect Tiebreaker Logic

The Problem: When multiple nodes would save the same number of nodes, the problem asks for the smallest indexed node. Some implementations might return the first one found or the last one found instead.

Wrong Implementation:

# WRONG: Returns last one found when tied
if component_size >= max_saved_nodes:  # Should be > not >=
    best_node_to_remove = infected_node

Correct Implementation:

# RIGHT: Properly handles ties by choosing smaller index
if (component_size > max_saved_nodes or 
    (component_size == max_saved_nodes and infected_node < best_node_to_remove)):
    best_node_to_remove = infected_node
    max_saved_nodes = component_size

Complete Solution to Avoid These Pitfalls:

The provided solution correctly handles all these cases by:

1. Using Union-Find to identify connected components without modifying the graph
2. Counting infected nodes per component to identify which can be saved
3. Only considering components with exactly one infected node
4. Having a fallback for when no component can be saved
5. Properly handling the tiebreaker by checking node indices

The key insight is that we can only prevent a component's infection if it has exactly one initially infected node, and we remove that specific node from the initial set.