C++で行列（2次元配列）の累積和（プレフィックスサム）を計算する方法

この記事では、整数値を格納した2次元配列 mat[][] が与えられたとき、その累積和行列（プレフィックスサム行列）を求める問題について解説します。

累積和行列とは？

累積和行列とは、行列の各要素が「その位置より上側と左側にあるすべての要素の合計」を表す行列のことです。数式で表すと以下のようになります。

prefixSum[i][j] = mat[i][j] + mat[i-1][j] + ... + mat[0][j] + mat[i][j-1] + ... + mat[i][0]

つまり、(i, j) の値は、左上の要素 (0, 0) から (i, j) までの長方形領域内の全要素の総和に等しくなります。

具体例

入力: arr = [
    [4   6   1]
    [5   7   2]
    [3   8   9]
]
出力: [
    [4    10   11]
    [9    22   25]
    [12   33   45]
]

例えば出力の 22 は、入力行列の左上からその位置までの範囲（4 + 6 + 5 + 7）の合計になっています。

解き方：単純なアプローチと効率的なアプローチ

最も単純な方法は、各位置 (i, j) に対して、そこまでのすべての要素を毎回走査して合計を求めることです。しかし、この方法では計算量が O(R² × C²) となり、行列サイズが大きくなると非効率です。

そこで有効なのが、すでに計算済みの累積和の値を再利用する漸化式を使う方法です。これにより O(R × C) の計算量で求められます。

一般式

prefixSum[i][j] = prefixSum[i-1][j] + prefixSum[i][j-1] - prefixSum[i-1][j-1] + a[i][j]

ポイントは「上のセル」と「左のセル」の累積和を足し合わせると、左上のセルが二重にカウントされてしまうため、それを1回引くという点です。

特別なケース（端の処理）

i = 0 かつ j = 0 の場合:
prefixSum[i][j] = a[i][j]

i = 0 かつ j > 0 の場合（1行目）:
prefixSum[i][j] = prefixSum[i][j-1] + a[i][j]

i > 0 かつ j = 0 の場合（1列目）:
prefixSum[i][j] = prefixSum[i-1][j] + a[i][j]

C++での実装例

以下のコードは、上記のアルゴリズムを実装したものです。

#include <iostream>
using namespace std;
#define R 3
#define C 3

void printPrefixSum(int a[][C]) {
    int prefixSum[R][C];
    // 左上の要素
    prefixSum[0][0] = a[0][0];
    // 1行目の処理
    for (int i = 1; i < C; i++)
        prefixSum[0][i] = prefixSum[0][i - 1] + a[0][i];
    // 1列目の処理
    for (int i = 1; i < R; i++)
        prefixSum[i][0] = prefixSum[i - 1][0] + a[i][0];
    // 残りの要素を一般式で計算
    for (int i = 1; i < R; i++) {
        for (int j = 1; j < C; j++)
            prefixSum[i][j] = prefixSum[i - 1][j] + prefixSum[i][j - 1]
                            - prefixSum[i - 1][j - 1] + a[i][j];
    }
    // 結果の出力
    for (int i = 0; i < R; i++) {
        for (int j = 0; j < C; j++)
            cout << prefixSum[i][j] << "\t";
        cout << endl;
    }
}

int main() {
    int mat[R][C] = {
        { 1, 2, 3},
        { 4, 5, 6},
        { 7, 8, 9}
    };
    cout << "The prefix Sum Matrix is :\n";
    printPrefixSum(mat);
    return 0;
}

実行結果

The prefix Sum Matrix is :
1    3    6
5    12   21
12   27   45

まとめ

累積和行列は、動的計画法の考え方を応用した基本的かつ重要なテクニックです。各要素の計算に直前の結果を再利用することで、計算量を大幅に削減できます。この手法は、部分行列の合計を高速に求めたい場面（競技プログラミングや画像処理など）で広く活用されています。