Kotlinで二次元累積和と二次元区間和を求める方法

二次元累積和と二次元区間和とは、行列の一部分の要素の和を高速に求めることができるアルゴリズムです。この記事では、Kotlinで二次元累積和と二次元区間和を求める方法について説明します。

二次元累積和とは

二次元累積和とは、行列の左上から任意の位置までの部分行列の和を事前に計算しておくことで、任意の区間の和を高速に求めることができるテクニックです。

具体的には、以下のような手順で計算します。

• まず、s[0][0] を 0 とします。
• 次に、s[i][0] を i 行目の要素の和とします。つまり、s[i][0] = s[i - 1][0] + a[i - 1][0] となります。
• 同様に、s[0][j] を j 列目の要素の和とします。つまり、s[0][j] = s[0][j - 1] + a[0][j - 1] となります。
• そして、s[i][j] を左上から (i, j) の位置までの部分行列の和とします。つまり、s[i][j] = s[i - 1][j] + s[i][j - 1] - s[i - 1][j - 1] + a[i - 1][j - 1] となります。この式は、上の行と左の列の和を足して、重複した部分を引いて、現在の要素を足すことで求められます。

このようにして、s[i][j] をすべて計算することで、二次元累積和が完成します。

二次元区間和とは

二次元区間和とは、行列の任意の区間 (l, r) の要素の和を求めることです。ここで、区間 (l, r) は左上の位置 l = (a, b) と右下の位置 r = (c, d) の組で表されます。

二次元累積和があれば、任意の区間 (l, r) の和は s[r][r] - s[l - 1][r] - s[r][l - 1] + s[l - 1][l - 1] で求めることができます。この式は、右下から左上に向かって四角形を切り取っていくイメージです。

Kotlinでの実装例

以下にKotlinで二次元累積和と二次元区間和を求めるコード例を示します。

import java.util.*

fun main() {
    // 入力を受け取る
    val sc = Scanner(System.`in`)
    val h = sc.nextInt()
    val w = sc.nextInt()
    val n = sc.nextInt()
    val a = Array(h) { IntArray(w) }
    for (i in 0 until h) {
        for (j in 0 until w) {
            a[i][j] = sc.nextInt()
        }
    }
    val abcd = Array(n) { IntArray(4) }
    for (i in 0 until n) {
        for (j in 0..3) {
            abcd[i][j] = sc.nextInt() - 1 // 0-indexedにする
        }
    }

    // 二次元累積和を計算する
    val s = Array(h + 1) { IntArray(w + 1) }
    s[0][0] = 0
    for (i in 1..h) {
        for (j in 1..w) {
            s[i][j] = s[i - 1][j] + s[i][j - 1] - s[i - 1][j - 1] + a[i - 1][j - 1]
        }
    }

    // 各ペアについての区間和を求める
    for (i in 0 until n) {
        val a = abcd[i][0]
        val b = abcd[i][1]
        val c = abcd[i][2]
        val d = abcd[i][3]
        val ans = s[c + 1][d + 1] - s[a][d + 1] - s[c + 1][b] + s[a][b] // S({a,b} , {c,d}) の値は s[c + 1][d + 1] - s[a][d + 1] - s[c + 1][b] + s[a][b] に等しい
        println(ans)
    }
}

参考

ソース: Bing との会話 2023/9/26
: 二次元累積和