LeetCode每日一题(928. Minimize Malware Spread II)

You are given a network of n nodes represented as an n x n adjacency matrix graph, where the ith node is directly connected to the jth node if graph[i][j] == 1.

Some nodes initial are initially infected by malware. Whenever two nodes are directly connected, and at least one of those two nodes is infected by malware, both nodes will be infected by malware. This spread of malware will continue until no more nodes can be infected in this manner.

Suppose M(initial) is the final number of nodes infected with malware in the entire network after the spread of malware stops.

We will remove exactly one node from initial, completely removing it and any connections from this node to any other node.

Return the node that, if removed, would minimize M(initial). If multiple nodes could be removed to minimize M(initial), return such a node with the smallest index.

Example 1:

Input: graph = [[1,1,0],[1,1,0],[0,0,1]], initial = [0,1]
Output: 0

Example 2:

Input: graph = [[1,1,0],[1,1,1],[0,1,1]], initial = [0,1]
Output: 1

Example 3:

Input: graph = [[1,1,0,0],[1,1,1,0],[0,1,1,1],[0,0,1,1]], initial = [0,1]
Output: 1

Constraints:

• n == graph.length
• n == graph[i].length
• 2 <= n <= 300
• graph[i][j] is 0 or 1.
• graph[i][j] == graph[j][i]
• graph[i][i] == 1
• 1 <= initial.length < n
• 0 <= initial[i] <= n - 1
• All the integers in initial are unique.

---

用所有不在 initial 里的 nodes 组成一个 graph， 然后把 graph 分成 1 个或多个互相不连接的 components, 统计每个 component 与多少个 initial 中的 nodes 相关联， 如果某个 component 只与 1 个 initial 中的 node 相关联， 那就说明我们在移除该 node 之后， 这个 component 中的所有 nodes 都会免于感染。

---

use std::collections::HashSet;

impl Solution {
    fn find(parents: &mut Vec<usize>, target: usize) -> usize {
        let parent = parents[target];
        if parent == target {
            return parents[target];
        }
        let p = Solution::find(parents, parent);
        parents[target] = p;
        p
    }

    fn union(parents: &mut Vec<usize>, counts: &mut Vec<i32>, a: usize, b: usize) {
        let parent_a = Solution::find(parents, a);
        let parent_b = Solution::find(parents, b);
        if parent_a == parent_b {
            return;
        }
        parents[parent_a] = parent_b;
        counts[parent_b] = counts[parent_a] + counts[parent_b];
    }
    pub fn min_malware_spread(graph: Vec<Vec<i32>>, initial: Vec<i32>) -> i32 {
        let injected: HashSet<usize> = initial.iter().map(|&v| v as usize).collect();
        let mut parents: Vec<usize> = (0..graph.len()).collect();
        let mut counts: Vec<i32> = vec![1; graph.len()];
        for i in 0..graph.len() {
            for j in i + 1..graph[0].len() {
                if graph[i][j] == 1 && !injected.contains(&i) && !injected.contains(&j) {
                    Solution::union(&mut parents, &mut counts, i, j);
                }
            }
        }
        let mut sources = vec![HashSet::new(); graph.len()];
        for i in initial.clone() {
            for j in 0..graph.len() {
                if injected.contains(&j) || graph[i as usize][j] == 0 {
                    continue;
                }
                let parent = Solution::find(&mut parents, j);
                sources[parent].insert(i);
            }
        }
        let mut contributions = vec![0; graph.len()];
        for i in initial {
            contributions[i as usize] = 1;
        }
        for (i, mut source) in sources.into_iter().enumerate() {
            if source.len() == 1 {
                let s: Vec<i32> = source.drain().collect();
                contributions[s[0] as usize] += counts[i];
            }
        }
        let mut node = 0;
        let mut max = 0;
        for (i, c) in contributions.into_iter().enumerate() {
            if c > max {
                max = c;
                node = i as i32;
            }
        }
        node
    }
}