Minimize Malware Spread Solutions in Go

package leetcode

import (
	"math"

	"github.com/halfrost/LeetCode-Go/template"
)

func minMalwareSpread(graph [][]int, initial []int) int {
	if len(initial) == 0 {
		return 0
	}
	uf, minIndex, count, countMap := template.UnionFind{}, 0, math.MinInt64, map[int]int{}
	uf.Init(len(graph))
	for i := range graph {
		for j := range graph[i] {
			if i == j {
				break
			}
			if graph[i][j] == 1 {
				uf.Union(i, j)
			}
		}
	}
	for i := 0; i < len(graph); i++ {
		countMap[uf.Find(i)]++
	}
	for _, v := range initial {
		tmp := countMap[uf.Find(v)]
		if count == tmp && minIndex > v {
			minIndex = v
		}
		if count < tmp {
			minIndex = v
			count = tmp
		}
	}
	return minIndex
}