Solution.ts

class UnionFind {
    p: number[];
    size: number[];
    constructor(n: number) {
        this.p = Array(n)
            .fill(0)
            .map((_, i) => i);
        this.size = Array(n).fill(1);
    }

    find(x: number): number {
        if (this.p[x] !== x) {
            this.p[x] = this.find(this.p[x]);
        }
        return this.p[x];
    }

    union(a: number, b: number): boolean {
        const [pa, pb] = [this.find(a), this.find(b)];
        if (pa === pb) {
            return false;
        }
        if (this.size[pa] > this.size[pb]) {
            this.p[pb] = pa;
            this.size[pa] += this.size[pb];
        } else {
            this.p[pa] = pb;
            this.size[pb] += this.size[pa];
        }
        return true;
    }

    getSize(root: number): number {
        return this.size[root];
    }
}

function minMalwareSpread(graph: number[][], initial: number[]): number {
    const n = graph.length;
    const uf = new UnionFind(n);
    for (let i = 0; i < n; ++i) {
        for (let j = i + 1; j < n; ++j) {
            graph[i][j] && uf.union(i, j);
        }
    }
    let [ans, mx] = [n, 0];
    const cnt: number[] = Array(n).fill(0);
    for (const x of initial) {
        ++cnt[uf.find(x)];
    }
    for (const x of initial) {
        const root = uf.find(x);
        if (cnt[root] === 1) {
            const sz = uf.getSize(root);
            if (sz > mx || (sz === mx && x < ans)) {
                [ans, mx] = [x, sz];
            }
        }
    }
    return ans === n ? Math.min(...initial) : ans;
}