UnionFind.kt

/* author: jeovazero */
/**
 * Linking by size + Path halving 
*/
class UnionFind(size: Int) {
    val roots = IntArray(size)
    val sizes = IntArray(size)
    var rootsSize = size - 1
    init {
        for(i in 0..(size - 1)) {
            roots[i] = i
            sizes[i] = 0
        }
    }
    fun find(node: Int): Int {
        var i = node
        while(roots[i] != i) {
            roots[i] = roots[roots[i]]
            i = roots[i]
        }
        return i
    }
    fun union(a: Int, b: Int) {
        val fa = find(a)
        val fb = find(b)

        if (fa == fb) return

        if (sizes[fa] < sizes[fb]) {
            roots[fa] = fb
            sizes[fb] += sizes[fa]
        } else {
            roots[fb] = fa
            sizes[fa] += sizes[fb]
        }
        rootsSize--
    }
}