928. Minimize Malware Spread II

Problem

(This problem is the same as Minimize Malware Spread, with the differences bolded.)

In a network of nodes, each node i is directly connected to another node j if and only 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 one node from the initial list, 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,0],[0,0,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

Note:

1. 1 < graph.length = graph[0].length <= 300
2. 0 <= graph[i][j] == graph[j][i] <= 1
3. graph[i][i] = 1
4. 1 <= initial.length < graph.length
5. 0 <= initial[i] < graph.length

Problem Summary

(This problem is the same as Minimize Malware Spread, with the differences shown in bold.) In a network of nodes, each node i can be directly connected to another node j only when graph[i][j] = 1. Some nodes initial are initially infected by malware. As long as two nodes are directly connected and at least one of them 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 can remove one node from the initial list, completely removing that node and any connections from that node to any other node. Return the node that, if removed, would minimize M(initial). If multiple nodes satisfy the condition, return the node with the smallest index.

Notes:

• 1 < graph.length = graph[0].length <= 300
• 0 <= graph[i][j] == graph[j][i] <= 1
• graph[i][i] = 1
• 1 <= initial.length < graph.length
• 0 <= initial[i] < graph.length

Solution Ideas

• This problem is an enhanced version of problem 924. Given a graph describing relationships between nodes, if two nodes are connected, the malware will infect all nodes connected to them. Now, if we want to completely and thoroughly remove one infected node to reduce infections as much as possible, which node should be removed? If multiple nodes can reduce the number of infections, prioritize removing the node with the smaller index. The input and output requirements of this problem are exactly the same as problem 924. The difference is that problem 924 essentially asks to turn one infected node into an uninfected node, while this problem completely deletes one infected node and all edges connected to it.
• This problem tests union-find. Of course, DFS can also be used to solve this problem. The union-find approach is as follows: first remove all infected nodes, then merge all connected components into one node each. Because the nodes in a connected set are either all infected or all uninfected, each set can be considered as a whole. Then count the number of infected nodes directly adjacent to each set. For a set:
  1. If the number of directly adjacent infected nodes is 0, it definitely will not be infected, so ignore this case;
  2. If the number of directly adjacent infected nodes is 1, then after deleting that infected node, the entire connected component can avoid being infected. This is the case we are looking for;
  3. If the number of directly adjacent infected nodes is greater than or equal to 2, then no matter which infected node is deleted, the connected component will still be infected, so ignore this case;
• Therefore, we only need to find, among all connected components of the second type (connected components with exactly 1 directly adjacent infected node), the connected component with the largest number of nodes. The infected node adjacent to it is the node we should delete; if multiple connected components have the same number of nodes, then find the corresponding infected node with the smallest index.

Code

package leetcode

import (
	"math"

	"github.com/halfrost/leetcode-go/template"
)

func minMalwareSpread2(graph [][]int, initial []int) int {
	if len(initial) == 0 {
		return 0
	}
	uf, minIndex, count, countMap, malwareMap, infectMap := template.UnionFind{}, initial[0], math.MinInt64, map[int]int{}, map[int]int{}, map[int]map[int]int{}
	for _, v := range initial {
		malwareMap[v]++
	}
	uf.Init(len(graph))
	for i := range graph {
		for j := range graph[i] {
			if i == j {
				break
			}
			if graph[i][j] == 1 && malwareMap[i] == 0 && malwareMap[j] == 0 {
				uf.Union(i, j)
			}
		}
	}
	for i := 0; i < len(graph); i++ {
		countMap[uf.Find(i)]++
	}
	// Record the number of directly adjacent malware nodes for each set
	for _, i := range initial {
		for j := 0; j < len(graph); j++ {
			if malwareMap[j] == 0 && graph[i][j] == 1 {
				p := uf.Find(j)
				if _, ok := infectMap[p]; ok {
					infectMap[p][i] = i
				} else {
					tmp := map[int]int{}
					tmp[i] = i
					infectMap[p] = tmp
				}
			}
		}
	}
	// Select the malware node with the smallest index
	for _, v := range initial {
		minIndex = min(minIndex, v)
	}
	for i, v := range infectMap {
		// Find the ones connected to only one malware node
		if len(v) == 1 {
			tmp := countMap[uf.Find(i)]
			keys := []int{}
			for k := range v {
				keys = append(keys, k)
			}
			if count == tmp && minIndex > keys[0] {
				minIndex = keys[0]
			}
			if count < tmp {
				minIndex = keys[0]
				count = tmp
			}
		}
	}
	return minIndex
}