KotlinでUnion-Find木を実装

KotlinでUnion-Findを実装

AtCoderにてKotlinで実装をしているのですが、KotlinにはUnion-Find木が存在しないので自前で実装します。

class UnionFind(n: Int) {
    val par = Array(n){it}
    val rank = Array(n){0}

    fun find(x: Int): Int {
        if (par[x] == x) {
            return x
        } else {
            par[x] = find(par[x])
            return par[x]
        }
    }

    fun unite(x: Int, y: Int) {
        val px = find(x)
        val py = find(y)
        if (px == py) return
        if (rank[px] < rank[py]) {
            par[px] = py
        } else {
            par[py] = px
            if (rank[px] == rank[py]) {
                rank[px] ++
            }
        }
    }

    fun same(x: Int, y: Int): Boolean {
        return find(x) == find(y)
    }

}

参考: 蟻本