🗂

ABC129 D - Lamp のメモ

2023/10/28に公開

重実装系の問題を募集して教えてもらった問題です。延々とTLEを出していましたが最終的には解けました。
https://atcoder.jp/contests/abc129/tasks/abc129_d

考察

あるマスに明かりを設置したら、端または障害物があるところまで照らされるので、その照らされるマスの数の最大値を求める問題です。最近解いた問題とかなり似ています。
https://zenn.dev/dhirabayashi/articles/5eeb5494ac28e8

障害物がなければ計算するだけ、障害物があるならその位置を高速に求める必要があるという点が同じ。前回の問題は障害物の位置を二分探索で求めたので、同じ手が使えると思って実装しました。

長いのでコードは貼らずリンクだけにしますが、こちらはTLEとなりました。
https://atcoder.jp/contests/abc129/submissions/46907973

マス目の全探索を何度もやっているし、4方向それぞれに対して探索も行なっているし、ということでそりゃTLEになるだろと思うところですが、最初はなぜTLEになるのかよくわかっていませんでした。
HWが最大だと4 × 10^6通りもの数になるのですが、これがけっこう厳しい制約だと思っていなかったんですね…

いろいろ考えて、最終的には方針はそのままで探索の重複を除外して通すことができました。

import kotlin.math.max

fun main() {
    val (h, w) = readln().split(" ").map { it.toInt() }
    val s = listOf("") + List(h) {
        "." + readln()
    }

    solve(h, w, s)
}

fun solve(h: Int, w: Int, s: List<String>) {
    //var time = System.currentTimeMillis()

    // 障害物を二分探索するためのリスト
    val besides = Array(h + 1) { mutableListOf<Int>() }
    val verticals = Array(w + 1) { mutableListOf<Int>() }

    for(i in 1..h) {
        for(j in 1..w) {
            if(s[i][j] == '#') {
                besides[i].add(j)
                verticals[j].add(i)
            }
        }
    }

    //println("make: ${System.currentTimeMillis() - time}")

    var ans = 0

    // 探索
    //time = System.currentTimeMillis()

    for(i in 1..h) {
        for(j in 1..w) {
            if(s[i][j] == '#') {
                continue
            }

            var count = 0

            // 二分探索で障害物の場所を調べる
            // 横
            val besideIndex = lowerBound(besides[i], j)
            // 縦
            val verticalIndex = lowerBound(verticals[j], i)

            val besideCount = calcCount(besides[i], besideIndex, w)
            val verticalCount = calcCount(verticals[j], verticalIndex, h)
            val rawCount = besideCount + verticalCount
            count += if(rawCount == 0) {
                0
            } else {
                rawCount - 1
            }

            ans = max(ans, count)
        }
    }
    //println("bruteforce: ${System.currentTimeMillis() - time}")

    println(ans)
}

private fun lowerBound(list: List<Int>, item: Int): Int {
    val index = list.binarySearch(item)
    return -(index + 1)
}

private fun calcCount(barriers: List<Int>, barrierIndex: Int, end: Int): Int {
    if(barriers.isEmpty()) {
        return end
    }

    // 前(上)
    val front = if(barrierIndex == 0) {
        1
    } else {
        barriers[barrierIndex - 1] + 1
    }

    // 後(下)
    val back = if(barrierIndex == barriers.size) {
        end
    } else {
        barriers[barrierIndex] - 1
    }

    return back - front + 1
}

https://atcoder.jp/contests/abc129/submissions/46956849

まず二分探索をするためのデータを作りますが、TLEコードではこれを何度も実行していたのを一回に減らしました。
また、各マスについて上下左右の4方向それぞれに対して二分探索をしていましたが、これはよく考えたら横と縦の二回でいいとわかったのでその分も減らしました。
各マスに対応する値も、二次元配列で管理すれば単一の数値でいいので、Pairなどは使わないようにしました。
二分探索で、その探索対象の値がどの位置に挿入できるのかわかるので、その値が0なら前に障害物はないとわかりますし、0でないならその一つ前の値が障害物の位置とわかります。同様に配列の末尾の添字なら後ろに障害物はなく、そうでないなら一つ後ろの値が障害物の位置、とわかるのでそれをもとに照らされる範囲を計算できます。

最終的には通りましたが、

それでもけっこうギリギリです。4 × 10^6はかなり重いと認識を改めます…

別解法

二分探索を使わず、各マスが照らされる回数を単純に数えて記録する方法でも解けると教えてもらったのでそれもやってみました。

ある行について、横に連続する空きマスの個数は照らされる個数と等しいので、それを各行と列に対して計算する、障害物があったらリセットしてまた別途計算する、をやっていけばいいと。たしかに…シンプルなのに気づきませんでした。

import kotlin.math.max

fun main() {
    val (h, w) = readln().split(" ").map { it.toInt() }
    val s = Array(h) {
        readln()
    }

    // 縦用
    val vMutableIntMap = mutableMapOf<Int, MutableInt>()

    // マスごとの個数管理用
    val counts = Array(h) { Array<Point?>(w) { null } }

    for(i in 0 until h) {
        var beside = MutableInt()
        for(j in 0 until w) {
            if(s[i][j] == '#') {
                beside = MutableInt()
                vMutableIntMap[j] = MutableInt()
                continue
            }

            // 横
            beside.increment()

            // 縦
            vMutableIntMap[j] = vMutableIntMap.getOrDefault(j, MutableInt()).increment()

            counts[i][j] = Point(beside, vMutableIntMap[j]!!)
        }
    }

    var ans = 0
    for(i in 0 until h) {
        for(j in 0 until w) {
            ans = max(ans, counts[i][j]?.calc() ?: 0)
        }
    }
    println(ans)
}

class MutableInt(private var value: Int = 0) {
    fun increment(): MutableInt {
        value++
        return this
    }
    fun toInt() = value
}

data class Point(
    val i: MutableInt,
    val j: MutableInt,
) {
    fun calc() = i.toInt() + j.toInt() - 1
}

https://atcoder.jp/contests/abc129/submissions/46985346

空きマスの個数を数えていって、障害物が見つかったら個数が確定しますが、それより前のマスも一通りその値に更新する必要があります。一通り走査して書き換えていったら遅くて間に合わないので、ここでは可変なオブジェクトを使い回すことで対応しました。参照をコピーしているだけなので、書き換えたら全てに変更が反映されます。なので連続する空きマスの分だけインスタンスを作って引き回せばいいです。
横と縦それぞれに対して全探索すると無駄なので、縦については参照をMapで管理しています。

一通り値が求まったら、それの中で最大になる値を探索します。全探索が2回発生するのでちょっと嫌だなと思ったのですが、とりあえず大丈夫でした。

これでもまだちょっと遅いですが、最初の解法よりはだいぶ速いですね。

最初の解法より速い理由は二分探索をしなくて済んでいるからだと思います。二分探索は爆速と雑に認識していましたが、とはいえ回数が増えるとまあまあ時間かかるのですね。
(執筆時間: 29分34秒)

Discussion