924. Minimize Malware Spread #

Problem #

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. 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.

Note that if a node was removed from the initial list of infected nodes, it may still be infected later as a result of the malware spread.

Example 1:

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

Example 2:

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

Example 3:

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

Note:

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

Problem Summary #

In a network of nodes, only when graph[i][j] = 1 can each node i be directly connected to another node j. 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 malware spread will continue until no more nodes can be infected in this way. Suppose M(initial) is the final number of nodes infected by malware in the entire network after the malware stops spreading. We can remove one node from the initial list. If removing this node minimizes M(initial), return that node. If multiple nodes satisfy the condition, return the node with the smallest index. Note that if a node has been removed from the infected-node list initial, it may still be infected later due to malware spread.

Note:

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 #

Given a graph of relationships between nodes, if two nodes are connected, then the malware will infect all nodes in that connected component. Now, if we want to remove one infected node to reduce the infection 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.
At first glance, this problem is about Union-Find. Use the connectivity relationships between nodes to union() all the nodes given in the problem, then count how many nodes are in each set one by one. Finally, scan the initial array once and select the node in this array with the smaller index and whose set contains more nodes; this node is the final answer.

Code #


package leetcode

import (
	"math"

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

func minMalwareSpread(graph [][]int, initial []int) int {
	if len(initial) == 0 {
		return 0
	}
	uf, maxLen, maxIdx, uniqInitials, compMap := template.UnionFindCount{}, math.MinInt32, -1, map[int]int{}, map[int][]int{}
	uf.Init(len(graph))
	for i := range graph {
		for j := i + 1; j < len(graph); j++ {
			if graph[i][j] == 1 {
				uf.Union(i, j)
			}
		}
	}
	for _, i := range initial {
		compMap[uf.Find(i)] = append(compMap[uf.Find(i)], i)
	}
	for _, v := range compMap {
		if len(v) == 1 {
			uniqInitials[v[0]] = v[0]
		}
	}
	if len(uniqInitials) == 0 {
		smallestIdx := initial[0]
		for _, i := range initial {
			if i < smallestIdx {
				smallestIdx = i
			}
		}
		return smallestIdx
	}
	for _, i := range initial {
		if _, ok := uniqInitials[i]; ok {
			size := uf.Count()[uf.Find(i)]
			if maxLen < size {
				maxLen, maxIdx = size, i
			} else if maxLen == size {
				if i < maxIdx {
					maxIdx = i
				}
			}
		}
	}
	return maxIdx
}