package g0901_1000.s0924_minimize_malware_spread

// #Hard #Array #Depth_First_Search #Breadth_First_Search #Matrix #Union_Find
// #2023_04_17_Time_791_ms_(100.00%)_Space_59.7_MB_(100.00%)

class Solution {
    private lateinit var size: IntArray
    private lateinit var par: IntArray

    fun minMalwareSpread(graph: Array<IntArray>, initial: IntArray): Int {
        size = IntArray(graph.size)
        par = IntArray(graph.size)
        for (i in graph.indices) {
            size[i] = 1
            par[i] = i
        }
        // create groups
        for (i in graph.indices) {
            for (j in graph[0].indices) {
                if (graph[i][j] == 1) {
                    val p1 = find(i)
                    val p2 = find(j)
                    merge(p1, p2)
                }
            }
        }
        // no of infected in group
        val infected = IntArray(graph.size)
        for (e in initial) {
            val p = find(e)
            infected[p]++
        }
        var currSize = -1
        var ans = -1
        for (e in initial) {
            val p = find(e)
            if (infected[p] == 1 && size[p] >= currSize) {
                ans = if (size[p] > currSize) {
                    e
                } else {
                    ans.coerceAtMost(e)
                }
                currSize = size[p]
            }
        }
        // all groups have more than 1 infected node then return min value from initial
        if (ans == -1) {
            ans = initial[0]
            for (j in initial) {
                ans = ans.coerceAtMost(j)
            }
        }
        return ans
    }

    private fun merge(p1: Int, p2: Int) {
        if (p1 != p2) {
            if (size[p1] > size[p2]) {
                par[p2] = p1
                size[p1] += size[p2]
            } else {
                par[p1] = p2
                size[p2] += size[p1]
            }
        }
    }

    private fun find(u: Int): Int {
        if (par[u] == u) {
            return u
        }
        par[u] = find(par[u])
        return par[u]
    }
}
