LeetCode 928: Minimize Malware Spread II

Problem Restatement

We are given an undirected graph with n nodes.

The graph is represented as an adjacency matrix:

graph[i][j] == 1

means node i is directly connected to node j.

Some nodes are initially infected by malware. These nodes are listed in:

initial

Malware spreads through edges. If an infected node is connected to an uninfected node, the uninfected node also becomes infected. This continues until no more nodes can be infected.

We must remove exactly one node from initial. Unlike LeetCode 924, this removed node is completely removed from the graph, including all of its edges.

Return the node whose removal minimizes the final number of infected nodes.

If several nodes give the same minimum result, return the smallest node index.

The official statement defines an n x n adjacency matrix, removes one node from initial completely, and asks for the node minimizing the final infection count, with smallest index as the tie-breaker.

Input and Output

Function shape:

class Solution:
    def minMalwareSpread(self, graph: list[list[int]], initial: list[int]) -> int:
        ...

Examples

Example 1:

graph = [
    [1, 1, 0],
    [1, 1, 0],
    [0, 0, 1],
]
initial = [0, 1]

Nodes 0 and 1 are connected.

If we remove node 0, node 1 is still infected, but node 0 is gone.

Final infected count is 1.

If we remove node 1, node 0 is still infected, but node 1 is gone.

Final infected count is also 1.

Both removals are equally good, so we return the smaller index:

0

Example 2:

graph = [
    [1, 1, 0],
    [1, 1, 1],
    [0, 1, 1],
]
initial = [0, 1]

If we remove node 0, node 1 remains infected and spreads to node 2.

Final infected nodes are:

{1, 2}

If we remove node 1, node 0 remains infected, but node 2 is no longer connected through node 1.

Final infected nodes are:

{0}

So removing 1 is better.

Answer:

1

Example 3:

graph = [
    [1, 1, 0, 0],
    [1, 1, 1, 0],
    [0, 1, 1, 1],
    [0, 0, 1, 1],
]
initial = [0, 1]

If we remove node 0, node 1 can still spread through the chain to nodes 2 and 3.

If we remove node 1, node 0 remains infected, but nodes 2 and 3 are separated from the infection.

Answer:

1

These match the official examples for problem 928.

First Thought: Try Each Removal

The constraints are small enough to support a direct simulation.

The graph has at most:

n = 300

For each infected node removed in initial:

1. Pretend this node no longer exists.
2. Start malware spread from every other node in initial.
3. Count how many nodes become infected.
4. Choose the removed node that gives the smallest count.

This is simple and reliable.

The key difference from problem 924 is important: the removed node is not merely made uninfected. It is completely removed with all its connections.

Algorithm

First sort initial.

initial.sort()

Sorting handles the tie-breaker naturally. If two removals produce the same final infected count, we keep the first one we found, which is the smaller index.

For each candidate removed in initial:

1. Create a seen array.
2. Mark removed as already seen so traversal never enters it.
3. Add every other initially infected node to a queue.
4. Run BFS through the graph.
5. Count how many nodes are infected.
6. Update the answer if this count is smaller.

Return the best candidate.

Walkthrough

Use:

graph = [
    [1, 1, 0],
    [1, 1, 1],
    [0, 1, 1],
]
initial = [0, 1]

Try removing 0.

The remaining infected node is 1.

BFS starts from:

[1]

Node 1 connects to node 2, so node 2 becomes infected.

Final infected set:

{1, 2}

Count:

2

Try removing 1.

The remaining infected node is 0.

BFS starts from:

[0]

Node 0 cannot reach node 2 without node 1, because node 1 was removed.

Final infected set:

{0}

Count:

1

So the best removal is:

1

Correctness

For a fixed candidate removed, the algorithm marks removed as unavailable before the spread starts. Therefore, BFS never counts it and never travels through it. This exactly models the statement requirement that the node and its connections are completely removed.

The queue is initialized with every initially infected node except removed. These are exactly the remaining infection sources after the removal.

BFS then visits every node reachable from those infection sources without passing through removed. Malware spreads exactly along graph edges until no new reachable node can be infected, so BFS computes exactly the final infected set for that candidate.

The algorithm repeats this exact simulation for every possible node that can be removed from initial. Therefore, it evaluates every legal choice.

It stores the candidate with the smallest final infected count. Since initial is sorted and the answer is only updated when a strictly smaller count is found, ties keep the smaller index. This matches the required tie-breaker.

So the returned node is exactly the node whose removal minimizes the final malware spread.

Complexity

Let:

n = len(graph)
m = len(initial)

Since n <= 300, this direct simulation is acceptable.

Implementation

from collections import deque

class Solution:
    def minMalwareSpread(self, graph: list[list[int]], initial: list[int]) -> int:
        n = len(graph)
        initial.sort()

        best_node = initial[0]
        best_count = n + 1

        for removed in initial:
            infected_count = self._spread_count(graph, initial, removed)

            if infected_count < best_count:
                best_count = infected_count
                best_node = removed

        return best_node

    def _spread_count(
        self,
        graph: list[list[int]],
        initial: list[int],
        removed: int,
    ) -> int:
        n = len(graph)
        seen = [False] * n
        queue = deque()

        seen[removed] = True

        for node in initial:
            if node != removed:
                seen[node] = True
                queue.append(node)

        infected_count = 0

        while queue:
            node = queue.popleft()
            infected_count += 1

            for nei in range(n):
                if not seen[nei] and graph[node][nei] == 1:
                    seen[nei] = True
                    queue.append(nei)

        return infected_count

Code Explanation

We sort the initially infected nodes:

initial.sort()

This makes tie-breaking easy.

The first node is the default answer:

best_node = initial[0]
best_count = n + 1

Then we test every possible removal:

for removed in initial:
    infected_count = self._spread_count(graph, initial, removed)

If this candidate creates fewer final infections, we update the answer:

if infected_count < best_count:
    best_count = infected_count
    best_node = removed

We do not update on equal counts, because sorted order already gives the smaller index first.

Inside _spread_count, we block the removed node:

seen[removed] = True

This prevents BFS from entering it or using it as a bridge.

Then we add all remaining initially infected nodes to the queue:

for node in initial:
    if node != removed:
        seen[node] = True
        queue.append(node)

The BFS spreads infection through connected nodes:

for nei in range(n):
    if not seen[nei] and graph[node][nei] == 1:
        seen[nei] = True
        queue.append(nei)

The final infected_count is the value of M(initial) after removing that candidate.

Testing

def run_tests():
    s = Solution()

    assert s.minMalwareSpread(
        [
            [1, 1, 0],
            [1, 1, 0],
            [0, 0, 1],
        ],
        [0, 1],
    ) == 0

    assert s.minMalwareSpread(
        [
            [1, 1, 0],
            [1, 1, 1],
            [0, 1, 1],
        ],
        [0, 1],
    ) == 1

    assert s.minMalwareSpread(
        [
            [1, 1, 0, 0],
            [1, 1, 1, 0],
            [0, 1, 1, 1],
            [0, 0, 1, 1],
        ],
        [0, 1],
    ) == 1

    assert s.minMalwareSpread(
        [
            [1, 0, 0],
            [0, 1, 0],
            [0, 0, 1],
        ],
        [2],
    ) == 2

    assert s.minMalwareSpread(
        [
            [1, 1, 1],
            [1, 1, 1],
            [1, 1, 1],
        ],
        [2, 1],
    ) == 1

    print("all tests passed")

run_tests()