Minimize Malware Spread — DSA · Engineering Playbook

Problem#

Given an undirected graph (adjacency matrix graph) and a list of initially infected nodes initial, remove exactly one node from initial to minimize the final number of infected nodes (infection spreads through connected components). Return the chosen node; on tie, return the smallest index.

Examples#

graph = [[1,1,0],[1,1,0],[0,0,1]], initial = [0,1] → 0.
graph = [[1,0,0],[0,1,0],[0,0,1]], initial = [0,2] → 0.

Constraints#

n == graph.length, 2 <= n <= 300, 1 <= initial.length <= n.

Hints#

Hint 1

Build components with DSU. A component containing one initial node is “saved” if we remove that node.

Hint 2

A component with >= 2 initial nodes can’t be saved by removing only one — its infection persists.

Approach#

DSU all edges. For each initial node, count how many initial nodes share its component. The savable components are those with exactly one initial node — removing that node saves componentSize infections. Pick the initial node giving the largest saving, breaking ties by smallest index. If no component is single-seeded, return the smallest initial node.

Solution#

class Solution {
    vector<int> parent, sz;
    int find(int x) { while (parent[x] != x) { parent[x] = parent[parent[x]]; x = parent[x]; } return x; }
    void unite(int a, int b) {
        a = find(a); b = find(b); if (a == b) return;
        if (sz[a] < sz[b]) swap(a, b);
        parent[b] = a; sz[a] += sz[b];
    }
public:
    int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {
        int n = graph.size();
        parent.resize(n); sz.assign(n, 1);
        iota(parent.begin(), parent.end(), 0);
        for (int i = 0; i < n; ++i)
            for (int j = i + 1; j < n; ++j)
                if (graph[i][j]) unite(i, j);
        unordered_map<int,int> initInRoot;
        for (int v : initial) ++initInRoot[find(v)];
        sort(initial.begin(), initial.end());
        int best = initial[0], bestSave = -1;
        for (int v : initial) {
            int r = find(v);
            int save = (initInRoot[r] == 1) ? sz[r] : 0;
            if (save > bestSave) { bestSave = save; best = v; }
        }
        return best;
    }
};

import "sort"

func minMalwareSpread(graph [][]int, initial []int) int {
    n := len(graph)
    parent := make([]int, n)
    sz := make([]int, n)
    for i := 0; i < n; i++ {
        parent[i] = i
        sz[i] = 1
    }
    var find func(int) int
    find = func(x int) int {
        for parent[x] != x {
            parent[x] = parent[parent[x]]
            x = parent[x]
        }
        return x
    }
    unite := func(a, b int) {
        a, b = find(a), find(b)
        if a == b {
            return
        }
        if sz[a] < sz[b] {
            a, b = b, a
        }
        parent[b] = a
        sz[a] += sz[b]
    }
    for i := 0; i < n; i++ {
        for j := i + 1; j < n; j++ {
            if graph[i][j] == 1 {
                unite(i, j)
            }
        }
    }
    initInRoot := map[int]int{}
    for _, v := range initial {
        initInRoot[find(v)]++
    }
    sort.Ints(initial)
    best, bestSave := initial[0], -1
    for _, v := range initial {
        r := find(v)
        save := 0
        if initInRoot[r] == 1 {
            save = sz[r]
        }
        if save > bestSave {
            bestSave = save
            best = v
        }
    }
    return best
}

from collections import defaultdict
from typing import List

class Solution:
    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:
        n = len(graph)
        parent = list(range(n))
        sz = [1] * n

        def find(x: int) -> int:
            while parent[x] != x:
                parent[x] = parent[parent[x]]
                x = parent[x]
            return x

        def unite(a: int, b: int) -> None:
            a, b = find(a), find(b)
            if a == b:
                return
            if sz[a] < sz[b]:
                a, b = b, a
            parent[b] = a
            sz[a] += sz[b]

        for i in range(n):
            for j in range(i + 1, n):
                if graph[i][j] == 1:
                    unite(i, j)

        init_in_root: dict = defaultdict(int)
        for v in initial:
            init_in_root[find(v)] += 1

        initial = sorted(initial)
        best, best_save = initial[0], -1
        for v in initial:
            r = find(v)
            save = sz[r] if init_in_root[r] == 1 else 0
            if save > best_save:
                best_save = save
                best = v
        return best

function minMalwareSpread(graph, initial) {
    const n = graph.length;
    const parent = Array.from({ length: n }, (_, i) => i);
    const sz = new Array(n).fill(1);

    const find = (x) => {
        while (parent[x] !== x) {
            parent[x] = parent[parent[x]];
            x = parent[x];
        }
        return x;
    };
    const unite = (a, b) => {
        a = find(a); b = find(b);
        if (a === b) return;
        if (sz[a] < sz[b]) [a, b] = [b, a];
        parent[b] = a;
        sz[a] += sz[b];
    };

    for (let i = 0; i < n; i++) {
        for (let j = i + 1; j < n; j++) {
            if (graph[i][j] === 1) unite(i, j);
        }
    }

    const initInRoot = new Map();
    for (const v of initial) {
        const r = find(v);
        initInRoot.set(r, (initInRoot.get(r) ?? 0) + 1);
    }

    const sorted = initial.slice().sort((a, b) => a - b);
    let best = sorted[0], bestSave = -1;
    for (const v of sorted) {
        const r = find(v);
        const save = initInRoot.get(r) === 1 ? sz[r] : 0;
        if (save > bestSave) {
            bestSave = save;
            best = v;
        }
    }
    return best;
}

class Solution {
    private int[] parent, sz;

    public int minMalwareSpread(int[][] graph, int[] initial) {
        int n = graph.length;
        parent = new int[n];
        sz = new int[n];
        for (int i = 0; i < n; i++) {
            parent[i] = i;
            sz[i] = 1;
        }
        for (int i = 0; i < n; i++) {
            for (int j = i + 1; j < n; j++) {
                if (graph[i][j] == 1) unite(i, j);
            }
        }
        Map<Integer, Integer> initInRoot = new HashMap<>();
        for (int v : initial) {
            int r = find(v);
            initInRoot.merge(r, 1, Integer::sum);
        }
        Arrays.sort(initial);
        int best = initial[0], bestSave = -1;
        for (int v : initial) {
            int r = find(v);
            int save = (initInRoot.get(r) == 1) ? sz[r] : 0;
            if (save > bestSave) {
                bestSave = save;
                best = v;
            }
        }
        return best;
    }

    private int find(int x) {
        while (parent[x] != x) {
            parent[x] = parent[parent[x]];
            x = parent[x];
        }
        return x;
    }

    private void unite(int a, int b) {
        a = find(a);
        b = find(b);
        if (a == b) return;
        if (sz[a] < sz[b]) {
            int t = a; a = b; b = t;
        }
        parent[b] = a;
        sz[a] += sz[b];
    }
}

function minMalwareSpread(graph: number[][], initial: number[]): number {
    const n = graph.length;
    const parent: number[] = Array.from({ length: n }, (_, i) => i);
    const sz: number[] = new Array(n).fill(1);

    const find = (x: number): number => {
        while (parent[x] !== x) {
            parent[x] = parent[parent[x]];
            x = parent[x];
        }
        return x;
    };
    const unite = (a: number, b: number): void => {
        a = find(a); b = find(b);
        if (a === b) return;
        if (sz[a] < sz[b]) [a, b] = [b, a];
        parent[b] = a;
        sz[a] += sz[b];
    };

    for (let i = 0; i < n; i++) {
        for (let j = i + 1; j < n; j++) {
            if (graph[i][j] === 1) unite(i, j);
        }
    }

    const initInRoot: Map<number, number> = new Map();
    for (const v of initial) {
        const r = find(v);
        initInRoot.set(r, (initInRoot.get(r) ?? 0) + 1);
    }

    const sorted = initial.slice().sort((a, b) => a - b);
    let best = sorted[0], bestSave = -1;
    for (const v of sorted) {
        const r = find(v);
        const save = initInRoot.get(r) === 1 ? sz[r] : 0;
        if (save > bestSave) {
            bestSave = save;
            best = v;
        }
    }
    return best;
}

Editorial#

The crucial observation: a component with two or more initial infected nodes still gets fully infected after removing any one of them. Only components with exactly one initial node give a real saving — and that saving equals the component’s size. Sorting initial ascending and breaking ties with strict > makes the smallest-index tiebreak automatic.

Complexity#

Time: O(N^2 * α(N)) for the adjacency-matrix scan.
Space: O(N).

Concept revision#

Number of Connected Components in an Undirected Graph