Update solution 924

leetcode/0924.Minimize-Malware-Spread/924. Minimize Malware Spread.go

@@ -10,30 +10,43 @@ 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, 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 := range graph[i] {
-			if i == j {
-				break
-			}
+		for j := i + 1; j < len(graph); j++ {
 			if graph[i][j] == 1 {
 				uf.Union(i, j)
 			}
 		}
 	}
-	for i := 0; i < len(graph); i++ {
-		countMap[uf.Find(i)]++
+	for _, i := range initial {
+		compMap[uf.Find(i)] = append(compMap[uf.Find(i)], i)
 	}
-	for _, v := range initial {
-		tmp := countMap[uf.Find(v)]
-		if count == tmp && minIndex > v {
-			minIndex = v
-		}
-		if count < tmp {
-			minIndex = v
-			count = tmp
+	for _, v := range compMap {
+		if len(v) == 1 {
+			uniqInitials[v[0]] = v[0]
 		}
 	}
-	return minIndex
+	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
 }