累積和のお勉強

累積和

前処理をすることで、配列の区間の総和を$O(1)$で求めることができます。

1次元の累積和

prefix sum

前からの累積和を考えましょう。
$s[i]$を配列$a$の$[0, i)$の総和と定義します。

\begin{align*}
s[0] &= 0\\
s[1] &= s[0] + a[0]\\
s[2] &= s[1] + a[1]\\
\vdots
\end{align*}

という観察から$s[i] = s[i - 1] + a[i - 1]$が分かります。

prefix sum

int n;
cin >> n;
vector<int> a(n);
for (int i = 0; i < n; i++) cin >> a[i];

vector<int> s(n + 1);
for (int i = 1; i <= n; i++) s[i] = s[i - 1] + a[i - 1];

// 配列aの区間[l, r)の総和(l <= r)
// [l, r) = [0, r) - [0, l)
int l, r;
cin >> l >> r;
cout << s[r] - s[l] << endl;

suffix sum

$s[i]$を$[i, n)$の和と定義します。

\begin{align*}
s[n] &= 0\\
s[n - 1] &= s[n] + a[n - 1]\\
s[n - 2] &= s[n - 1] + a[n - 2]\\
\vdots
\end{align*}

という観察から$s[i] = s[i + 1] + a[i]$が分かります。

suffix sum

for (int i = n - 1; i >= 0; i--) s[i] = s[i + 1] + a[i];

// [l, r) = [l, n) - [r, n)
cout << s[l] - s[r] << endl;

例題：AtCoder ABC334 C-Socks 2

2次元の累積和

2次元配列の長方形の区間の総和を求めましょう。
2次元のsuffix sumも各次元ごとに累積和をとればできます。

回答クエリに関しては包除原理を用いますが、分からなくても多分大丈夫です。後で解説します。

2次元累積和

int h, w;
cin >> h >> w;
vector<vector<int>> a(h, vector<int>(w));
for (int i = 0; i < h; i++)
    for (int j = 0; j < w; j++) cin >> a[i][j];

vector<vector<int>> s1(h + 1, vector(w + 1));
auto s2 = s1;
// aに対して各行で1次元の累積和をとる
for (int i = 0; i <= h; i++)
    for (int j = 0; j <= w; j++)
        s1[i][j] = s1[i][j - 1] + a[i][j - 1];
// s1に対して各列で1次元の累積和をとる
for (int i = 0; i <= w; i++)
    for (int j = 0; j <= w; j++) 
        s2[i][j] = s2[i - 1][j] + s1[i - 1][j];
    
// 包除原理
// [a, b) * [c, d)の長方形領域の和
int a, b, c, d;
cin >> a >> b >> c >> d;
cout << s2[c][d] - s2[a][d] - s2[c][b] + s2[a][b] << endl;

ちなみに、$s1, s2$のように別々に配列を作らなくてもin-placeな感じに実装できます。ただ、$a$を書き換えのはよろしくない気がするので$s1, s2$を一つにまとめて$s$とします。

vector<vector<int>> s(h + 1, vector<int>(w + 1));
for (int i = 1; i <= h; i++) 
    for (int j = 1; j <= w; j++)
        s[i][j] = a[i - 1][j - 1];

for (int i = 1; i <= h; i++)
    for (int j = 1; j <= w; j++)
        s[i][j] += s[i][j - 1];
        
for (int i = 1; i <= h; i++)
    for (int j = 1; j <= w; j++)
        s[i][j] += s[i - 1][j];

例題：AtCoder 競技プログラミングの鉄則 演習問題集 A08-Two Dimensional Sum

3次元の累積和

2次元と同様に各次元ごとに累積和をとればよいです。

3次元の累積和

int h, w, d;
cin >> h >> w >> d;
vector a(h, vector<vector<int>>(w, vector<int>(d)));
for (int i = 0; i < h; i++)
    for (int j = 0; j < w; j++)
        for (int k = 0; k < d; k++)
            cin >> a[i][j][k];
            
vector s(h + 1, vector<vector<int>>(w + 1, vector<int>(d + 1)));
for (int i = 1; i <= h; i++)
    for (int j = 1; j <= w; j++)
        for (int k = 1; k <= d; k++)
            s[i][j][k] = a[i - 1][j - 1][k - 1];
            
for (int i = 1; i <= h; i++)
    for (int j = 1; j <= w; j++)
        for (int k = 1; k <= d; k++)
            s[i][j][k] += s[i][j][k - 1];

for (int i = 1; i <= h; i++)
    for (int j = 1; j <= w; j++)
        for (int k = 1; k <= d; k++)
            s[i][j][k] += s[i][j - 1][k];
        
for (int i = 1; i <= h; i++)
    for (int j = 1; j <= w; j++)
        for (int k = 1; k <= d; k++)
            s[i][j][k] += s[i - 1][j][k];
            
// 回答クエリ:[lx, rx) * [ly, ry) * [lz, rz)
auto sum_z = [&](int rx, int ry, int lz, int rz) {
    return s[rx][ry][rz] - s[rx][ry][lz];
};

auto sum_yz = [&](int rx, int ly, int ry, int lz, int rz) {
    return sum_z(rx, ry, lz, rz) - sum_z(rx, ly, lz, rz);
};

auto sum_xyz = [&](int lx, int rx, int ly, int ry, int lz, int rz) {
    return sum_yz(rx, ly, ry, lz, rz) - sum_yz(lx, ly, ry, lz, rz);
};

cout << sum_xyz(lx, rx, ly, ry, lz, rz) << endl;

例題：AtCoder ABC366 D-Cuboid Sum Query

包除原理

さて、回答クエリの部分での包除原理について3次元の場合で解説します。
$[lx, rx) * [ly, ry) * [lz, rz)$の形が計算しにくいので$[0, ?)$の形に直したいです。

まず、$y, z$を固定して考えると

$$[lx, rx) * [ly, ry) * [lz, rz) = \underline{[0, rx) * [ly, ry) * [lz, rz)} - [0, lx) * [ly, ry) * [lz, rz)$$

となります。$[0, rx)、[0, lx)$の形はうれしいです。

次に、$z$を固定して考えると

$$\underline{[0, rx) * [ly, ry) * [lz, rz)} = [0, rx) * [0, ry) * [lz, rz) - [0, rx) *[0, ly) * [lz, rz)$$

となります。残すは$[lz, rz)$だけですが1次元と同様に考えれば$[lz, rz] = [0, rz) - [0, lz)$なので

$$[0, rx) * [0, ry) * [lz, rz) = [0, rx] * [0, ry) * [0, rz) - [0, rx] * [0, ry) * [0, lz)$$

となります。

これで全て$[0, ?)$の形にできるので事前に計算した累積和の配列$s$によって計算できます。たとえば、$[0, rx] * [0, ry) * [0, rz)$は$s[rx][ry][rz]$に対応しています。

おわりに

よくわからないよーという人は以下の解説を見るといいと思います。

・累積和を何も考えずに書けるようにする！

・添字や境界条件で毎回バグらせてる人かわいそう

・２次元累積和の４つの形と８つの構築方法、４種類の長方形和クエリ、そして閃光のごとく現れた１人の天才こと私です。オススメ実装もあります！

・3次元の累積和で挙げた例題の解説放送