2次元累積和
(data_structure/cumulative_sum_2d.hpp)

• View this file on GitHub
• Last update: 2023-04-16 00:50:22+09:00
• Include: #include "data_structure/cumulative_sum_2d.hpp"

使い方

1. ある2次元配列の矩形和を求める場合
   • コンストラクタ引数に2次元配列を入力
   • sum(lx, ly, rx, ry)で $ [lx, rx) \times [ly, ry) $ の和を返す
2. 2次元imos法のように使う場合
   • コンストラクタ引数に入力した値をh, wとすると、$[0, h) \times [0, w)$ への加算ができる(加算しうる半開区間の右端の値を入力する)
   • imos(lx, ly, rx, ry, z)で $ [lx, rx) \times [ly, ry) $ に $ z $ を加算
   • build()で構築
   • rui.get(x, y)でアクセス

Verified with

• verify/aoj_dsl/aoj_dsl_5_b.test.cpp

Code

#pragma once

template <class T> struct CumulativeSum2D {
    int h, w;
    std::vector<std::vector<T>> seg;

    CumulativeSum2D() = default;

    CumulativeSum2D(int h, int w) : h(h), w(w), seg(h + 1, std::vector<T>(w + 1, T(0))) {}

    CumulativeSum2D(std::vector<std::vector<T>> &a) {
        h = (int)a.size();
        assert(h > 0);
        w = (int)a[0].size();
        seg.assign(h + 1, std::vector<T>(w + 1, T(0)));
        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w; j++) {
                seg[i + 1][j + 1] = a[i][j];
            }
        }
        for (int i = 0; i < h + 1; i++) {
            for (int j = 0; j < w; j++) {
                seg[i][j + 1] += seg[i][j];
            }
        }
        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w + 1; j++) {
                seg[i + 1][j] += seg[i][j];
            }
        }
    }

    // [lx, rx) x [ly, ry)
    T sum(int lx, int ly, int rx, int ry) {
        assert(0 <= lx and lx <= rx and rx <= h);
        assert(0 <= ly and ly <= ry and ry <= w);
        return (seg[rx][ry] - seg[lx][ry] - seg[rx][ly] + seg[lx][ly]);
    }

    // (i, j) \in [lx, rx) x [ly, ry) seg[i][j] += z;
    void imos(int lx, int ly, int rx, int ry, T z = T(1)) {
        assert(0 <= lx and lx <= rx and rx <= h);
        assert(0 <= ly and ly <= ry and ry <= w);
        seg[lx][ly] += z;
        seg[lx][ry] -= z;
        seg[rx][ly] -= z;
        seg[rx][ry] += z;
    }

    void build() {
        for (int i = 0; i < h + 1; i++) {
            for (int j = 0; j < w; j++) {
                seg[i][j + 1] += seg[i][j];
            }
        }
        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w + 1; j++) {
                seg[i + 1][j] += seg[i][j];
            }
        }
    }

    // return A[x][y]
    T get(int x, int y) const {
        assert(0 <= x and x < h and 0 <= y and y < w);
        return seg[x][y];
    }

    // output
    friend std::ostream &operator<<(std::ostream &os, const CumulativeSum2D &A) {
        os << "h = " << A.h << ", w = " << A.w << "\n";
        for (int i = 0; i <= A.h; i++) {
            for (int j = 0; j <= A.w; j++) {
                os << A.seg[i][j] << " \n"[j == A.w];
            }
        }
        return os;
    }
};

#line 2 "data_structure/cumulative_sum_2d.hpp"

template <class T> struct CumulativeSum2D {
    int h, w;
    std::vector<std::vector<T>> seg;

    CumulativeSum2D() = default;

    CumulativeSum2D(int h, int w) : h(h), w(w), seg(h + 1, std::vector<T>(w + 1, T(0))) {}

    CumulativeSum2D(std::vector<std::vector<T>> &a) {
        h = (int)a.size();
        assert(h > 0);
        w = (int)a[0].size();
        seg.assign(h + 1, std::vector<T>(w + 1, T(0)));
        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w; j++) {
                seg[i + 1][j + 1] = a[i][j];
            }
        }
        for (int i = 0; i < h + 1; i++) {
            for (int j = 0; j < w; j++) {
                seg[i][j + 1] += seg[i][j];
            }
        }
        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w + 1; j++) {
                seg[i + 1][j] += seg[i][j];
            }
        }
    }

    // [lx, rx) x [ly, ry)
    T sum(int lx, int ly, int rx, int ry) {
        assert(0 <= lx and lx <= rx and rx <= h);
        assert(0 <= ly and ly <= ry and ry <= w);
        return (seg[rx][ry] - seg[lx][ry] - seg[rx][ly] + seg[lx][ly]);
    }

    // (i, j) \in [lx, rx) x [ly, ry) seg[i][j] += z;
    void imos(int lx, int ly, int rx, int ry, T z = T(1)) {
        assert(0 <= lx and lx <= rx and rx <= h);
        assert(0 <= ly and ly <= ry and ry <= w);
        seg[lx][ly] += z;
        seg[lx][ry] -= z;
        seg[rx][ly] -= z;
        seg[rx][ry] += z;
    }

    void build() {
        for (int i = 0; i < h + 1; i++) {
            for (int j = 0; j < w; j++) {
                seg[i][j + 1] += seg[i][j];
            }
        }
        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w + 1; j++) {
                seg[i + 1][j] += seg[i][j];
            }
        }
    }

    // return A[x][y]
    T get(int x, int y) const {
        assert(0 <= x and x < h and 0 <= y and y < w);
        return seg[x][y];
    }

    // output
    friend std::ostream &operator<<(std::ostream &os, const CumulativeSum2D &A) {
        os << "h = " << A.h << ", w = " << A.w << "\n";
        for (int i = 0; i <= A.h; i++) {
            for (int j = 0; j <= A.w; j++) {
                os << A.seg[i][j] << " \n"[j == A.w];
            }
        }
        return os;
    }
};

Back to top page