AtCoder B,C問題をKotlinで解こう - ABC267

当記事ではAtCoder、ABCのB問題ならびにC問題（時々D問題も）のKotlinでの解法を超初心者向けに詳細に解説します。

B - Split?

やりたいこと

7列の中から 3列(列A,B,C, A<B<C） を選び出す。全ての3列の組み合わせ について、以下の条件を全て満たすものが1つ以上存在すれば”スプリット"となる。

• ピン1が倒れている。
• 列Aに1つ以上の倒れていないピンがある。
• 列Bのピンが全て倒れている。
• 列Cに1つ以上の倒れていないピンがある。

入力値の取得

    // 入力値の取得
    val s = readLine()!!

サンプルコード

main.kt

fun main(args: Array<String>) {
    println(if (getAns()) "Yes" else "No")
}

fun getAns(): Boolean {
    // 入力値の取得
    val s = readLine()!!

    // ピンの位置を定義
    val pins = listOf(setOf(7), setOf(4), setOf(2, 8), setOf(1, 5), setOf(3, 9), setOf(6), setOf(10))

    // ピン1が立っているか倒れているかを確認
    if (s[0] == '1') {
        // ピン1が立っているとスプリットの条件を達成できない
        return false
    }

    // 列Aを処理
    for (a in (0..pins.lastIndex - 2)) {
        // 列Aで立っているピンの有無を確認
        if (pins[a].all { s[it - 1] == '0' }) {
            // 列Aで全てのピンが倒れているので条件を満たせない
            continue
        }
        for (b in (a + 1..pins.lastIndex - 1)) {
            // 列Bで立っているピンの有無を確認
            if (pins[b].any { s[it - 1] == '1' }) {
                // 列Bに立っているピンが存在するので条件を満たせない
                continue
            }
            // 列Cで立っているピンの有無を確認
            for (c in (b + 1..pins.lastIndex)) {
                if (pins[c].all { s[it - 1] == '0' }) {
                    // 列Cで全てのピンが倒れているので条件を満たせない
                    continue
                }
                // 全ての条件を満たせた
                return true
            }
        }
    }
    // 全ての条件を満たせる組み合わせが一つもなかった
    return false
}

C - Index × A(Continuous ver.)

やりたいこと

$a_i$ から始まる長さMの連続部分列を式に当てはめて計算した結果を仮に $f(i)$ とする。
仮にMを4とすると、
$f(1) = (a_1 * 1) + (a_2 * 2) + (a_3 * 3) + (a_4 * 4)$
$f(2) = (a_2 * 1) + (a_3 * 2) + (a_4 * 3) + (a_5 * 4)$

となる。
$f(2) = (a_1 * 0) + (a_2 * 1) + (a_3 * 2) + (a_4 * 3) + (a_5 * 4)$
としてみると、

$f(i) = f(i-1) - (a_{i-1} + a_{i} + a_{i+1} + a_{i+2}) + (a_{i+3} * 4)$

となることがわかる。

$a_i$ から始まる長さMの部分列の合計を $s(i)$ とすると。
$f(i) = f(i-1) - s(i-1) + (a_{i+m-1} * m)$ となる。

$f(i)$ の計算結果を保持して、$f(i+1)$ の計算に利用しよう。
$s(i)$ をあらかじめまとめて計算して、処理を高速化しよう。

入力値の取得

    // 入力値の取得
    val (n, m) = readLine()!!.split(" ").map { it.toInt() }
    val a = readLine()!!.split(" ").map { it.toInt() }

サンプルコード

main.kt

fun main(args: Array<String>) {
    // 入力値の取得
    val (n, m) = readLine()!!.split(" ").map { it.toInt() }
    val a = readLine()!!.split(" ").map { it.toInt() }

    // s(i) ~ (a[i]..a[i+m-1]) の合計を保持する配列
    val subTotal = LongArray(n)
    // s(0) は a の先頭からm個の要素を合計する。
    subTotal[0] = a.take(m).map { it.toLong() }.sum()
    // s(1)以降は一つ前の要素から差分を計算することで高速に求められる
    for (i in 1..a.size - m) {
        subTotal[i] = subTotal[i - 1] - a[i - 1] + a[i + m - 1]
    }

    // f(i) ~ (a[i]*1,a[i+1]*2 .. a[i+m-1]*m)の合計を保持する配列
    val total = LongArray(n)
    // f(0) は aの先頭の要素からm個の要素を抜き出し、インデックス値+1を乗ずる。
    total[0] = (0 until m).map { a[it].toLong() * (it + 1) }.sum()
    // 答え（最大値）を保持する変数
    var ans = total[0]
    // f(1)以降は一つ前の要素から差分を計算することで高速に求められる
    for (i in 1..a.size - m) {
        total[i] = total[i - 1] - subTotal[i - 1] + a[i + m - 1].toLong() * m
        // 最大値の更新
        ans = Math.max(ans, total[i])
    }
    // 答えの出力
    println(ans)
}