Solution.go

func minMalwareSpread(graph [][]int, initial []int) int {
	n := len(graph)
	p := make([]int, n)
	size := make([]int, n)
	clean := make([]bool, n)
	for i := 0; i < n; i++ {
		p[i], size[i], clean[i] = i, 1, true
	}
	for _, i := range initial {
		clean[i] = false
	}

	var find func(x int) int
	find = func(x int) int {
		if p[x] != x {
			p[x] = find(p[x])
		}
		return p[x]
	}
	union := func(a, b int) {
		pa, pb := find(a), find(b)
		if pa != pb {
			size[pb] += size[pa]
			p[pa] = pb
		}
	}

	for i := 0; i < n; i++ {
		if !clean[i] {
			continue
		}
		for j := i + 1; j < n; j++ {
			if clean[j] && graph[i][j] == 1 {
				union(i, j)
			}
		}
	}
	cnt := make([]int, n)
	mp := make(map[int]map[int]bool)
	for _, i := range initial {
		s := make(map[int]bool)
		for j := 0; j < n; j++ {
			if clean[j] && graph[i][j] == 1 {
				s[find(j)] = true
			}
		}
		for root, _ := range s {
			cnt[root]++
		}
		mp[i] = s
	}
	mx, ans := -1, 0
	for i, s := range mp {
		t := 0
		for root, _ := range s {
			if cnt[root] == 1 {
				t += size[root]
			}
		}
		if mx < t || (mx == t && i < ans) {
			mx, ans = t, i
		}
	}
	return ans
}