🔥

ABC226 C - Martial artistのメモ

2022/11/24に公開

AtCoder

競プロ

グラフ

tech

目についたC問題を解いてみるの回です。

考察

これは技同士の依存関係をグラフとして表せば解けそうだなとすぐに気づきました。ある技が、その技を覚えるために必要な別の技を参照する有向グラフです。

グラフは二次元配列を使って表現することができるというのは以前学びました。

以前学んだのは無向グラフとして相互に参照するようにしていましたが、今回のは有向グラフとして一方通行のグラフとして表せます。
入力を素直に受け取ると自ずとそうなるので、けっこう親切な問題かもしれません。

グラフで持つのは技名を表す数値だけなので、技名と習得に必要な時間とをマッピングする情報を別途持ちます。

まず入力をもとにそれらの情報を生成します。生成したら、Nから探索を開始して依存する技を再帰的に参照していって習得に必要な時間を足していって合計値を求めればそれが答えです。いわゆるDFSです。
前提となる技が何もなしで習得できる技が必ずあるはずで、そこが再帰の終了条件となります。実際に技を覚える順番としてはそっちが出発点となるわけですが、実装上は逆から始めたほうが楽そうです。逆から始めれば習得に不要な技は自然と除外され、必要な分だけ探索できます。また、答えは単純な足し算で求まるため、逆から計算していっても答えは一致します。

この考え方だと絶対に必要な技だけ探索するので、最小値うんぬんは特に気にしなくていいはずです。
ただし、習得済の技を再度習得しようとする可能性があるため、そのガードは設けます。

なお、この問題は素直に考えても特にTLEにはならないと思います。

実装

var graph: MutableList<MutableList<Int>> = mutableListOf()

var table: IntArray = IntArray(0)

fun main() {
    Thread(null, {

        val n = readLine()!!.toInt()

        graph = MutableList(n+1) { mutableListOf<Int>() }

        table = IntArray(n+1)

        repeat(n) {
            val input = readLine()!!.split(" ").map { it.toInt() }
            val t = input[0]
            val k = input[1]
            val a = input.subList(2, 2 + k).toMutableList()
            val i = it + 1

            table[i] = t

            graph[i] = a
        }
        println(dfs(n, mutableSetOf()))
    }, "", 128 * 1024 * 1024).start()

}

fun dfs(x: Int, walked: MutableSet<Int>): Long {
    if(walked.contains(x)) {
        return 0L
    }

    walked.add(x)

    var sum = table[x].toLong()
    for(node in graph[x]) {
        sum += dfs(node, walked)
    }
    return sum
}

graphは技の依存関係を表すグラフで、前述の通り二次元配列で表せます。n+1確保することで1-indexedに補正します。
tableは技と習得に必要な時間をマッピングする配列です。これは連想配列を使ってもいいのですが、配列のほうが少し速いと思うので配列を使っています。1-indexedに補正するのはgraphと同じです。
これらはdfsから参照できるようにグローバルで定義していますが、ローカル変数にして引数で渡すのでも別によかったですね。初期化のために無駄な値を生成してしまっていますし。

dfsの実装は再帰がわからないと難しいと思いますが、再帰がわかっていればわかりやすい部類だと思います。walkedは探索済の技を再度探索しないためのガードです。
（変数名の命名はちょっと自信がないですが）
依存関係なしで習得できる技があった場合は、for文に入らないのでそこで再帰呼び出しが止まります。

あとは、メインスレッドで実行せずにスタックを拡張した別スレッドを生成して実行しています。メインスレッドで実行したらREになってしまったので。
スタックを拡張したスレッドで実行したら通ったので、スタックオーバーフローが発生していたと思われます。以前グラフを学んだときの問題もそうだったので、KotlinでDFSを実装するとかなりスタックオーバーフローになりやすいようですね…
今度からは必ずスタック拡張して実装しようと思います。

追記

あれこれやって実行時間を200msちょい削減しました。walkedもグルーバルにしてメソッド呼び出しじのオーバーヘッドを減らした、ListをArrayやIntArrayにしてオートボクシング/アンボクシングのコストをオーバーヘッドを減らした、あと入力時にBufferedReaderを使って改善した、などです。

var graph = Array(1) { IntArray(0) }

var table = IntArray(0)

var walked = mutableSetOf<Int>()

fun main() {
    Thread(null, {

        val n = readLine()!!.toInt()

        graph = Array(n+1) { IntArray(0) }

        table = IntArray(n+1)

        walked = mutableSetOf()

        val br = System.`in`.bufferedReader()
        repeat(n) {
            val input = br.readLine()!!.split(" ").map { s -> s.toInt() }
            val t = input[0]
            val k = input[1]
            val a = input.subList(2, 2 + k).toIntArray()
            val i = it + 1

            table[i] = t

            graph[i] = a
        }
        println(dfs(n))
    }, "", 128 * 1024 * 1024).start()
}

fun dfs(x: Int): Long {
    if(walked.contains(x)) {
        return 0L
    }

    walked.add(x)

    var sum = table[x].toLong()
    for(node in graph[x]) {
        sum += dfs(node)
    }
    return sum
}

考察

実装

追記

Discussion