AtCoder ABC293 D 解法メモ KotlinでUnion-Findしたい

Union-Findのお勉強をしました。

解いていた問題

atcoder.jp

解法メモ

• ロープの両端をノードとして、そこに最初からリンクが貼ってあるものとして考えればよい
• ロープの組を作っていく方法はいくつかあるけど、Union-Findでサクっと解きたい
  • Union-Findについては『問題解決力を鍛える！アルゴリズムとデータ構造』（大槻 兼資，秋葉 拓哉）｜講談社BOOK倶楽部で勉強しなおしました。 (前に一度読んでいたけれども全く記憶に残ってなかった・・・)

実装メモ

• 上記の書籍を元にKotlinで実装しなおしたのがこちら

private class UnionFind {
    val parentNode: MutableList<Int>
    val groupSize: MutableList<Int>
 
    constructor(n: Int) {
        parentNode = MutableList(n) { -1 }
        groupSize = MutableList(n) { 1 }
    }
 
    /**
     * 根を求める
     */
    fun root(x: Int): Int {
        return if (parentNode[x] == -1) {
            x // x が根の場合は x を返す
        } else {
            root(parentNode[x])
        }
    }
 
    /**
     * x と y が同じグループに属するかどうか
     * (根が一致するかどうか)
     */
    fun isSameGroup(x: Int, y: Int): Boolean {
        return root(x) == root(y)
    }
 
    /**
     * x を含むグループと y を含むグループとを併合する
     */
    fun unite(x: Int, y: Int): Boolean {
        // x, y をそれぞれ根まで移動する
        val xx = root(x)
        val yy = root(y)
 
        // すでに同じグループのときは何もしない
        if (xx == yy) return false
 
        // union by size (y 側のサイズが小さくなるようにする)
        if (groupSize[xx] < groupSize[yy]) {
            parentNode[xx] = yy
            groupSize[yy] = groupSize[yy] + groupSize[xx]
        } else {
            parentNode[yy] = xx
            groupSize[xx] = groupSize[xx] + groupSize[yy]
        }
        return true
    }
 
    /**
     * x を含むグループのサイズ
     */
    fun size(x: Int): Int {
        return groupSize[root(x)]
    }
 
 
    /**
     * このUnionFind内のグループの個数
     */
    fun groupCount(): Int {
        return parentNode.count { it == -1 }
    }
}

提出した回答

• Submission #39800100 - AtCoder Beginner Contest 293