7K12 blog

猫でも分かるアルゴリズム解説

NyaaLIB::DS_NyaaImos2D

二次元いもす法 矩形区間加算処理O(1)、矩形区間総和取得処理O(1)、更新処理O(HW)

#include <bits/stdc++.h>

namespace NyaaLIB
{
	/**
	* 二次元いもす法
	* 矩形区間加算処理O(1)、矩形区間総和取得処理O(1)、更新処理O(HW)
	**/

	template <class T = int64_t> struct DS_NyaaImos2D
	{
		int64_t ysz, xsz;
		std::vector<std::vector<T>> add;
		std::vector<std::vector<T>> sum;
		std::vector<std::vector<T>> val;
		// 半開区間[0,ysize)[0,xsize)のメモリ領域を確保する
		DS_NyaaImos2D(int64_t ysize, int64_t xsize) : ysz(ysize), xsz(xsize),
			add(ysize, vector<T>(xsize, 0)), val(ysize, vector<T>(xsize, 0)), sum(ysize + 1, vector<T>(xsize + 1, 0)) {}
		// 配列を入力としたときはメモリ確保と同時に累積和も作成する
		DS_NyaaImos2D(const std::vector<std::vector<T>>& vv) : ysz(int64_t(vv.size())), xsz(int64_t(vv[0].size())),
			add(ysz, vector<T>(xsz, 0)), val(ysz, vector<T>(xsz, 0)), sum(ysz + 1, vector<T>(xsz + 1, 0))
		{	// 0で初期化されたres[0][x]とres[y][0]も追加してsumを構築
			for (int64_t y = 1; y < ysz + 1; y++) for (int64_t x = 1; x < xsz + 1; x++)
			{
				sum[y][x] = sum[y - 1][x] + sum[y][x - 1] - sum[y - 1][x - 1] + vv[y - 1][x - 1];
			}
		}

		void Add(int64_t y, int64_t x, T z)
		{	// 二次元配列[y][x]に値zを加算
			add[y][x] += z;
			add[y][x + 1] -= z;
			add[y + 1][x] -= z;
			add[y + 1][x + 1] += z;
		}

		void Add(int64_t ymin, int64_t ymax, int64_t xmin, int64_t xmax, T z)
		{	// 閉区間[ymin, ymax],[xmin, xmax]に値zを加算
			add[ymin][xmin] += z;
			add[ymin][xmax + 1] -= z;
			add[ymax + 1][xmin] -= z;
			add[ymax + 1][xmax + 1] += z;
		}

		// 矩形範囲閉区間[ymin,ymax][xmin,xmax]の総和をO(1)で返す
		T Sum(int64_t ymin, int64_t ymax, int64_t xmin, int64_t xmax)
		{	// 累積和sumは元の座標と1ずつずれているので調整
			ymin++; xmin++; ymax++; xmax++;
			return sum[ymin - 1][xmin - 1] + sum[ymax][xmax] - sum[ymax][xmin - 1] - sum[ymin - 1][xmax];
		}

		decltype(val)& Update(void)
		{	// 更新処理
			for (int64_t y = 0; y < ysz; y++) for (int64_t x = 1; x < xsz; x++) add[y][x] += add[y][x - 1];
			for (int64_t y = 1; y < ysz; y++) for (int64_t x = 0; x < xsz; x++) add[y][x] += add[y - 1][x];
			for (int64_t y = 0; y < ysz; y++) for (int64_t x = 0; x < xsz; x++) val[y][x] += add[y][x];
			for (int64_t y = 0; y < ysz; y++) for (int64_t x = 0; x < xsz; x++) add[y][x] = 0;
			for (int64_t y = 1; y < ysz + 1; y++) for (int64_t x = 1; x < xsz + 1; x++)
			{
				sum[y][x] = sum[y - 1][x] + sum[y][x - 1] - sum[y - 1][x - 1] + val[y - 1][x - 1];
			}
			return val;
		}

		// 二次元配列[y][x]の値を返す
		T Val(int64_t y, int64_t x) { return val[y][x]; }
	};
}

実装例
https://atcoder.jp/contests/gigacode-2019/submissions/23839555