Minimize Malware Spread

Minimize Malware Spread is a classic Union-Find problem. Malware spreads across a network of connected computers; a few start infected. You may remove one initially-infected computer from the network before the spread happens, and you want to remove the one that saves the most machines.

The slow way first

The brute-force idea: for each candidate in initial, actually remove it, then run a full BFS/DFS flood from the remaining infected nodes and count how many get infected. Keep whichever removal leaves the fewest infected. That is correct but does a whole graph traversal per candidate — O(k * (n + edges)), and on a dense matrix that is O(k * n²).

The question to ask: which computers can possibly infect each other at all? Only computers in the same connected component. If we knew the components up front, we would not need to re-flood the graph for every candidate.

The idea: components, then count the infected per component

Build the connected components with Union-Find. Then the rule falls out:

• A component is only saveable if it contains exactly one infected node. Remove that node and the whole component stays clean.
• If a component has two or more infected nodes, removing just one leaves another to re-infect everything — that component is unsavable, so it contributes zero saved machines.

So: count infected nodes per component, then among components with exactly one infected node, remove the infected node belonging to the largest such component (ties → smallest index).

The key insight: a node only helps if it is the single point of infection for its component. Otherwise removing it is wasted.

Walk through it

Step through the animation. First we union the edges into components: A = {0, 1, 2}, B = {3, 4}, C = {5}. Component A has exactly one infected node (0) and size 3, so removing 0 saves 3 machines. Component B has two infected nodes (3 and 4), so it saves nothing. Only A qualifies, so the answer is 0.

Pseudocode

infected = set(initial)
union every pair of directly-connected computers   # build components
for each component, count how many infected nodes it holds
best, best_save = smallest index in initial, -1
for each candidate in sorted(initial):
    root = component of candidate
    if that component has exactly ONE infected node
       and its size > best_save:
        best, best_save = candidate, size of that component
return best

The Python solution

def min_malware_spread(graph, initial):
    inf = set(initial)
    uf = UnionFind(len(graph))
    for i in range(len(graph)):
        for j in range(i + 1, len(graph)):
            if graph[i][j]:
                uf.union(i, j)
    count = {}                       # root -> # infected
    for node in initial:
        count[uf.find(node)] = count.get(uf.find(node), 0) + 1
    best, best_save = min(initial), -1
    for node in sorted(initial):
        root = uf.find(node)
        if count[root] == 1 and uf.size(root) > best_save:
            best, best_save = node, uf.size(root)
    return best

• uf is a Union-Find that also tracks each component's size.
• The double loop unions every directly-connected pair, grouping the graph into connected components.
• count maps a component root → number of infected nodes in that component.
• best_save is the size of the best component found so far; we start it at -1 so any valid candidate wins.
• The decision line is count[root] == 1 — only a component with a single infected node can be saved by removing one node.
• Iterating over sorted(initial) means the first candidate at a given save-size wins, giving the smallest-index tie-break for free.

Complexity

The matrix scan is O(n²) no matter what, but Union-Find lets us answer every candidate from one pass plus near-constant-time find/union (the inverse-Ackermann a(n) factor), instead of re-flooding the graph k times.

When this pattern shows up

Practice