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
}