競プロノート

Appearance

累積和（Prefix Sum） ​

累積和は、配列の各位置までの和を前計算し、任意の区間の和をO(1)で求めるテクニック。

1D・2D・3Dに対応する。

特徴 ​

次元前処理区間和クエリ
1DO(N)O(1)
2DO(HW)O(1)
3DO(HWD)O(1)

添え字の基本ルール ​

このライブラリでは、すべての区間を 0-indexedの半開区間 で指定する。

次元区間指定意味
1D[l, r)l以上r未満
2D[xl, xr) × [yl, yr)長方形領域
3D[xl, xr) × [yl, yr) × [zl, zr)直方体領域

右端の値は含まれない。

例えば [1, 4) は添え字1, 2, 3を表す。

1D累積和 ​

考え方 ​

配列が次のとき、

text
添え字 i  0  1  2  3
値 a[i]   1  2  3  4

累積和は、

text
添え字 i  0  1  2  3   4
ps[i]     0  1  3  6  10

となる。

重要なのは、ps[i] が a[0] から a[i-1] までの和を表すこと。

cpp
ps[0] = 0;
ps[1] = a[0];
ps[2] = a[0] + a[1];
ps[3] = a[0] + a[1] + a[2];

つまり、長さNの配列に対して、累積和は長さN+1になる。

区間和の求め方 ​

[l, r) の和は、

cpp
ps[r] - ps[l]

で求められる。

例えば [1, 3) の場合、

text
a = [1, 2, 3, 4]
        └──┘
       対象区間

a[1] + a[2] = 2 + 3 = 5

累積和から求めると、

cpp
ps[3] - ps[1] = 6 - 1 = 5;

となる。

実装 ​

cpp

#pragma once
#include <bits/stdc++.h>
using namespace std;

template<class T = long long>
struct PrefixSum1D {
    int n = 0;
    vector<T> ps = {T(0)};
    bool built = false;

    PrefixSum1D() = default;
    explicit PrefixSum1D(int n) : n(n), ps(n + 1, T(0)) { assert(n >= 0); }
    explicit PrefixSum1D(const vector<T>& a) { build(a); }

    void add(int i, T x) {
        assert(!built && 0 <= i && i < n);
        ps[i + 1] += x;
    }
    void build() {
        assert(!built);
        for (int i = 0; i < n; i++) ps[i + 1] += ps[i];
        built = true;
    }
    void build(const vector<T>& a) {
        n = (int)a.size(); ps.assign(n + 1, T(0)); built = false;
        for (int i = 0; i < n; i++) add(i, a[i]);
        build();
    }
    T range_sum(int l, int r) const {
        assert(built && 0 <= l && l <= r && r <= n);
        return ps[r] - ps[l];
    }
    T range_sum_inclusive(int l, int r) const { return range_sum(l, r + 1); }
    T prefix(int r) const { return range_sum(0, r); }
    T suffix(int l) const { return range_sum(l, n); }
    T at(int i) const { return range_sum(i, i + 1); }

    // 全要素が非負の場合のみ使用する
    int lower_bound_prefix(T x) const {
        assert(built);
        auto it = lower_bound(ps.begin(), ps.end(), x);
        return it == ps.end() ? n + 1 : (int)(it - ps.begin());
    }
    // 累積和がx以下となる最大の右端。存在しなければ-1
    int upper_bound_prefix(T x) const {
        assert(built);
        return (int)(upper_bound(ps.begin(), ps.end(), x) - ps.begin()) - 1;
    }
};

使用例 ​

cpp
vector<ll> a = {1, 2, 3, 4};
PrefixSum1D<ll> ps(a);
操作対象結果
range_sum(0, 4)1+2+3+410
range_sum(1, 3)2+35
range_sum(2, 4)3+47
range_sum(1, 2)22
range_sum(2, 2)空区間0
prefix(3)1+2+36
suffix(2)3+47
at(2)a[2]3

閉区間で指定したい場合 ​

[l, r] のように右端も含めたい場合は、

cpp
ps.range_sum_inclusive(l, r);

を使用する。

例えば、

cpp
ps.range_sum(1, 3);           // 2+3 = 5
ps.range_sum_inclusive(1, 3); // 2+3+4 = 9

この2つの違いに注意する。

累積和上の二分探索 ​

配列の全要素が非負なら、累積和は単調非減少になる。

cpp
vector<ll> a = {1, 2, 3, 4};
PrefixSum1D<ll> ps(a);

// ps = {0, 1, 3, 6, 10}
操作意味結果
lower_bound_prefix(3)累積和が3以上になる最小の右端2
lower_bound_prefix(5)累積和が5以上になる最小の右端3
lower_bound_prefix(11)条件を満たす右端がない5
upper_bound_prefix(3)累積和が3以下となる最大の右端2

lower_bound_prefix() が見つからない場合は n+1、upper_bound_prefix() が見つからない場合は -1 を返す。

負の要素がある場合は、累積和が単調にならない可能性があるため、この二分探索は使用できない。

2D累積和 ​

2次元配列の長方形領域の和をO(1)で求める。

考え方 ​

次の配列を考える。

text
       y=0  y=1  y=2
x=0     1    2    3
x=1     4    5    6

累積和 ps[x][y] は、左上の (0, 0) から (x, y) の直前までの長方形の和を表す。

つまり、

ps[x][y] = Σ a[i][j] (0 <= i < x, 0 <= j < y)

となる。

累積和配列は各次元を1つ大きく確保する。

text
       y=0  y=1  y=2  y=3
x=0     0    0    0    0
x=1     0    1    3    6
x=2     0    5   12   21

例えば、

cpp
ps[2][2] = 1 + 2 + 4 + 5; // 12
ps[2][3] = 1 + 2 + 3 + 4 + 5 + 6; // 21

となる。

区間和の求め方 ​

長方形 [xl, xr) × [yl, yr) の和を求める。

左上の大きな領域から、不要な上側・左側を引き、重複して引いた部分を戻す。

cpp
ps[xr][yr]
- ps[xl][yr]
- ps[xr][yl]
+ ps[xl][yl]

実装 ​

cpp

#pragma once
#include <bits/stdc++.h>
using namespace std;

template<class T = long long>
struct PrefixSum2D {
    int h, w;
    vector<vector<T>> ps;
    bool built = false;

    PrefixSum2D(int h, int w) : h(h), w(w), ps(h + 1, vector<T>(w + 1, T(0))) { assert(h >= 0 && w >= 0); }
    explicit PrefixSum2D(const vector<vector<T>>& a)
        : PrefixSum2D((int)a.size(), a.empty() ? 0 : (int)a[0].size()) {
        for (int i = 0; i < h; i++) {
            assert((int)a[i].size() == w);
            for (int j = 0; j < w; j++) add(i, j, a[i][j]);
        }
        build();
    }
    void add(int x, int y, T val) {
        assert(!built && 0 <= x && x < h && 0 <= y && y < w);
        ps[x + 1][y + 1] += val;
    }
    void build() {
        assert(!built);
        for (int i = 1; i <= h; i++) for (int j = 1; j <= w; j++)
            ps[i][j] += ps[i - 1][j] + ps[i][j - 1] - ps[i - 1][j - 1];
        built = true;
    }
    // [xl,xr) × [yl,yr)
    T range_sum(int xl, int xr, int yl, int yr) const {
        assert(built && 0 <= xl && xl <= xr && xr <= h && 0 <= yl && yl <= yr && yr <= w);
        return ps[xr][yr] - ps[xl][yr] - ps[xr][yl] + ps[xl][yl];
    }
};

使用例 ​

cpp
vector<vector<ll>> a = {
    {1, 2, 3},
    {4, 5, 6}
};

PrefixSum2D<ll> ps(a);
操作対象結果
range_sum(0, 2, 0, 3)全体21
range_sum(0, 1, 0, 3)1行目6
range_sum(1, 2, 0, 3)2行目15
range_sum(0, 2, 1, 3)2+3+5+616
range_sum(1, 2, 1, 2)55

引数の順番に注意 ​

cpp
ps.range_sum(xl, xr, yl, yr);

引数は、x方向の左端・右端、y方向の左端・右端の順番。

cpp
ps.range_sum(0, 2, 1, 3);

なら、

  • x方向：0, 1行目
  • y方向：1, 2列目

が対象になる。

閉区間 [xl, xr] × [yl, yr] を指定したい場合は、

cpp
ps.range_sum(xl, xr + 1, yl, yr + 1);

とする。

3D累積和 ​

3次元配列の直方体領域の和をO(1)で求める。

考え方 ​

2D累積和を3次元に拡張したもの。

ps[x][y][z] は、

[0, x) × [0, y) × [0, z)

に含まれる値の総和を表す。

元の配列が H × W × D の場合、累積和の論理的なサイズは、

(H+1) × (W+1) × (D+1)

となる。

実装内部ではメモリ効率のため1次元配列に格納しているが、考え方は同じ。

区間和の求め方 ​

直方体領域の区間和は、包除原理を使って8項の加減算で求める。

cpp
ps[xr][yr][zr]
- ps[xl][yr][zr]
- ps[xr][yl][zr]
- ps[xr][yr][zl]
+ ps[xl][yl][zr]
+ ps[xl][yr][zl]
+ ps[xr][yl][zl]
- ps[xl][yl][zl]

※ 上記は数式としての3次元配列表現。実装内部では1次元配列を使用している。

実装 ​

cpp

#pragma once
#include <bits/stdc++.h>
using namespace std;

template<class T = long long>
struct PrefixSum3D {
    int h, w, d;
    vector<T> ps;
    bool built = false;

    size_t id(int x, int y, int z) const {
        return ((size_t)x * ((size_t)w + 1) + y) * ((size_t)d + 1) + z;
    }

    PrefixSum3D(int h, int w, int d) : h(h), w(w), d(d) {
        assert(h >= 0 && w >= 0 && d >= 0);
        ps.assign(((size_t)h + 1) * ((size_t)w + 1) * ((size_t)d + 1), T(0));
    }
    explicit PrefixSum3D(const vector<vector<vector<T>>>& a)
        : PrefixSum3D((int)a.size(), a.empty() ? 0 : (int)a[0].size(), a.empty() || a[0].empty() ? 0 : (int)a[0][0].size()) {
        for (int i = 0; i < h; i++) {
            assert((int)a[i].size() == w);
            for (int j = 0; j < w; j++) {
                assert((int)a[i][j].size() == d);
                for (int k = 0; k < d; k++) add(i, j, k, a[i][j][k]);
            }
        }
        build();
    }
    void add(int x, int y, int z, T val) {
        assert(!built && 0 <= x && x < h && 0 <= y && y < w && 0 <= z && z < d);
        ps[id(x + 1, y + 1, z + 1)] += val;
    }
    void build() {
        assert(!built);
        for (int x = 1; x <= h; x++) for (int y = 0; y <= w; y++) for (int z = 0; z <= d; z++)
            ps[id(x, y, z)] += ps[id(x - 1, y, z)];
        for (int x = 0; x <= h; x++) for (int y = 1; y <= w; y++) for (int z = 0; z <= d; z++)
            ps[id(x, y, z)] += ps[id(x, y - 1, z)];
        for (int x = 0; x <= h; x++) for (int y = 0; y <= w; y++) for (int z = 1; z <= d; z++)
            ps[id(x, y, z)] += ps[id(x, y, z - 1)];
        built = true;
    }
    // [xl,xr) × [yl,yr) × [zl,zr)
    T range_sum(int xl, int xr, int yl, int yr, int zl, int zr) const {
        assert(built && 0 <= xl && xl <= xr && xr <= h && 0 <= yl && yl <= yr && yr <= w && 0 <= zl && zl <= zr && zr <= d);
        auto get = [&](int x, int y, int z) { return ps[id(x, y, z)]; };
        return get(xr, yr, zr) - get(xl, yr, zr) - get(xr, yl, zr) - get(xr, yr, zl)
             + get(xl, yl, zr) + get(xl, yr, zl) + get(xr, yl, zl) - get(xl, yl, zl);
    }
};

使用例 ​

cpp
PrefixSum3D<ll> ps(2, 2, 2);

ps.add(0, 0, 0, 5);
ps.add(0, 1, 1, 3);
ps.add(1, 1, 1, 7);
ps.build();

このとき、値がある座標は次の3つ。

座標値
(0,0,0)5
(0,1,1)3
(1,1,1)7

使用例：

cpp
ps.range_sum(0, 2, 0, 2, 0, 2); // 15
ps.range_sum(0, 1, 0, 2, 0, 2); // 8
ps.range_sum(1, 2, 1, 2, 1, 2); // 7

引数の順番に注意 ​

cpp
ps.range_sum(xl, xr, yl, yr, zl, zr);

各次元について、左端・右端の順番で指定する。

閉区間で指定したい場合は、

cpp
ps.range_sum(xl, xr + 1, yl, yr + 1, zl, zr + 1);

とする。

コンストラクタとbuildの違い ​

配列から構築する場合 ​

配列を渡すコンストラクタでは、内部で自動的にbuildされる。

cpp
PrefixSum1D<ll> ps1(a);
PrefixSum2D<ll> ps2(b);
PrefixSum3D<ll> ps3(c);

// そのままクエリを呼べる

サイズを指定して構築する場合 ​

add() で値を設定し、最後に build() を呼ぶ。

cpp
PrefixSum2D<ll> ps(h, w);
ps.add(0, 0, 5);
ps.add(1, 1, 3);
ps.build();

build() 後の値の追加・変更には対応していない。

よくある間違い ​

右端を含めてしまう ​

cpp
ps.range_sum(1, 3);

対象は a[1], a[2] であり、a[3] は含まれない。

累積和の添え字と元配列の添え字を混同する ​

cpp
ps[3]

は元配列の a[3] ではなく、a[0] + a[1] + a[2]。

2D・3Dで引数の順序を間違える ​

cpp
ps.range_sum(xl, xr, yl, yr);

であり、

cpp
ps.range_sum(xl, yl, xr, yr);

ではない。

buildを忘れる ​

サイズ指定コンストラクタを使用した場合は、build() が必要。

値を更新してしまう ​

累積和は静的な配列を対象とする。

値の更新が必要な場合は、BIT・セグメント木などを利用する。

関連 ​

  • いもす法
  • BIT
  • セグメント木
  • 包除原理