あるエンジニアのAtCoder奮闘記

東京都港区にあるアミフィアブル株式会社のエンジニアが、AtCoderで解いた問題について振り返ったりしていく会社公認のブログです。

AtCoder ABC 132 E - Hopscotch Addict (Kotlin)

irisruneです。どうにか自力ACできたのでE問題も取り上げることにしました。

問題

atcoder.jp

AtCoder上に類題がほぼないらしく、正解率が低めだったそうです。最短経路問題を理解しているならば、試行錯誤によって解法を思いつくのは割と難しくないかもしれません。

import java.util.*
import java.util.zip.Inflater

const val INF = 1000000000

class Node(val num: Int) {
    val adj = mutableListOf<Node>()
    var dMod = mutableListOf(INF, INF, INF)

    fun adjAdd(v: Node) {
        adj.add(v)
    }
}

fun main(args: Array<String>){
    val (n, m) = readLine()!!.split(" ").map(String::toInt)
    val graph = mutableListOf(Node(0))
    for (i in 1 until n + 1) graph.add(Node(i))
    for (j in 0 until m) {
        val (u, v) = readLine()!!.split(" ").map(String::toInt)
        graph[u].adjAdd(graph[v])
    }

    val (s, t) = readLine()!!.split(" ").map(String::toInt)
    graph[s].dMod[0] = 0
    var depth = 1
    var queue = ArrayDeque(listOf(graph[s]))
    do {
        val nextQueue = ArrayDeque<Node>()
        while (queue.size > 0) {
            val v = queue.poll()
            v.adj.forEach {
                if (it.num == t && depth % 3 == 0) {
                    println(depth / 3)
                    return
                }
                if (it.dMod[depth % 3] > depth) {
                    it.dMod[depth % 3] = depth
                    nextQueue.add(it)
                }
            }
        }
        queue = nextQueue
        depth++
    } while (queue.size > 0)
    println(-1)
}

まずベースとなるのは最短経路問題であり、辺の重みが1で固定されていることからDFSあるいはBFSを用いるのが無難です。

次にこの問題の本質である、3回の移動をワンセットとすることについて考えます。3回の移動を1セットにするということは始点からの距離を3で割った剰余が0,1,2それぞれの場合について扱いが異なるということになります。これを実現する方法は例えば、

  • 各頂点を3つに分けて辺を適切に張り直す(公式解説の方法)。
  • 始点から(始点を含む)各頂点への距離を3で割った剰余によって3つ用意する(上記提出コードの方法)。

などが考えられます。今回用いたのは後者の方法だったわけですが、メリットとしては実装が比較的簡単なことが挙げられます。一方デメリットとしては、BFS(あるいはDFS)の探索打ち切り(枝刈り?)条件が少々複雑になることが挙げられます。これはどういうことかというと、例えば始点からの距離1(剰余1)が既に記録されている頂点に距離3(剰余0)となる経路が見つかった場合は探索が続行されますが、距離4(剰余1)となる経路が見つかった場合は探索が打ち切られるといった具合です。

また、アルゴリズムの計算量が定数倍部分でよくないのか、DFSで実装するとTLE(さらには恐らく再帰が深すぎたためにREも)してしまったためにBFSで実装することとなりました。

雑記

  • 色々考えるに、個人的にはリスト内包表記が使えてなおかつ高速な言語があると嬉しい気がしてきました。
    • リスト内包表記自体は好みが分かれる面もあるようです。