ABC120 D Decayed BridgesをKotlinで - あるエンジニアのAtCoder奮闘記

irisruneです。今回もKotlinで解いた過去問についてですが、少し重い問題なので1問だけです。

橋を減らしていくのではなく逆に増やしていくという発想ができれば愚直に解くことはできると思います。そして時間内に解くのはUnion-Findを知らなくてもできます（できました）、でもD問題にしては実装も重い気がします。

class Bridge(val a : Int, val b : Int)
class Island(var group : Int)

fun moveGroup(s : MutableList<Int>, t : MutableList<Int>, lIsland : MutableList<Island>, group : Int){
    for(i in 0..(s.size - 1)){
        t.add(s[0])
        lIsland[s[0]].group = group
        s.removeAt(0)
    }
}

fun main(args: Array<String>){
    val (n, m) = readLine()!!.split(" ").map(String::toInt)
    var lIsland = mutableListOf<Island>()
    var lGroup : MutableList<MutableList<Int>> = mutableListOf()
    for(i in 0..(n - 1)){
        lIsland.add(Island(i))
        val iLGroup = mutableListOf(i)
        lGroup.add(iLGroup)
    }
    var lBridge = mutableListOf<Bridge>()
    for(i in 0..(m - 1)){
        val(a, b) = readLine()!!.split(" ").map(String::toInt)
        lBridge.add(Bridge(a - 1, b - 1))
    }
    var lRet = mutableListOf<Long>()
    lRet.add(n.toLong() * (n - 1).toLong() / 2L)
    for(i in (m - 1) downTo 1){
        val groupA = lIsland[lBridge[i].a].group
        val groupB = lIsland[lBridge[i].b].group
        if(groupA == groupB){
            lRet.add(lRet[(m - 1) - i])
        }else{
            lRet.add(lRet[(m - 1) - i] - (lGroup[groupA].size * lGroup[groupB].size))
            if(lGroup[groupA].size >= lGroup[groupB].size) moveGroup(lGroup[groupB], lGroup[groupA], lIsland, groupA)
            else moveGroup(lGroup[groupA], lGroup[groupB], lIsland, groupB)
        }
    }
    for(i in (m - 1) downTo 0){
        println(lRet[i])
    }
}

この問題の実装面での課題は解説にもある通り、

頂点 $u$ が属するグループと頂点 $v$ が属するグループを併合し、同じグループにする操作
頂点 $v$ が属するグループ番号を得る操作
頂点 $v$ が属するグループと同じグループに属する頂点数を得る操作

以上の操作をどのように実装するかになります。愚直に実装すると $O(N+M)$ 時間かかる操作が出てしまうので（よくわかってない）、全体で $O(M(N+M))$ 時間かかってしまい到底時間内には解けません。

今回自分はどのようにこれらの操作について実装したかというと、

頂点数の少ないグループの頂点を頂点数の多いグループに移動させる
- 全体で $O(N)$ 時間に。
頂点にグループ番号を持たせる
- 1操作あたり $O(1)$ 時間に。
グループごとにリストを作ってそこに頂点を入れてsizeを用いる
- 1操作あたり $O(1)$ 時間に。

以上のように実装することで全体で $O(N+M)$ 時間に抑えたので、無事ACできました。でもこの実装がUnion-Findかというと多分違うと思います。

今回のコードレビューは、

varを多用しすぎ。
一つの変数にIntを使うべき場所とLongを使うべき場所の両方が存在してしまい。toLong()を使う必要が出ている。
クラスの作りが非常に雑、~~こんなクラスならなくていい。~~

ざっと見た感じだとこのくらいですかね。~~多分もっと突っ込みどころありそうですね。~~

雑記

ACできたとは書きましたが、1800msくらいかかっていたので多分入力がボトルネックになっています。でも入力を実装するとブログ上にコード貼り付けるときに困るんですよね…
ABC120 C Unificationは愚直に解いてTLE出し続けました。他人の実行時間を見てから解けたので事実上自力ではないです。