Fenwick Sum

components/fenwick_sum.hpp · GitHub で見る

共通ページ

#include "lib/components/fenwick_sum.hpp"

FenwickSum<long long> bit(std::vector<long long>{3, 1, 4});
bit.add(1, 5);
auto sum = bit.prod(0, 3); // 13

区間は \([l,r)\)、添字は 0 始まりです。T{} が零元で、+=・-= が可換群の加減算になる型を使います。

操作	内容・返り値	計算量
FenwickSum() / FenwickSum(n)	空／零元が n 個の列を構築	空: \(O(1)\)、長さ指定: \(O(n)\)
FenwickSum(a) / build(a)	配列から構築／再構築	\(O(n)\)
size()	要素数	\(O(1)\)
add(i,x)	要素に x を加算	\(O(\log n)\)
add_pair(i,x,j,y)	2 点に加算。共通の祖先はまとめて更新	\(O(\log n)\)
prefix(r)	\([0,r)\) の和	\(O(\log n)\)
prod(l,r)	区間和。両端で共通する祖先の読み込みを省く	\(O(\log n)\)
prod_as<Accum>(l,r)	加減算を Accum 型で行い、区間和を返す	\(O(\log n)\)

prod_as は保持する値の型を変えず、取得時の加減算だけを指定した型で行います。木の各部分和は T、取得中の値と結果は Accum に収まる必要があります。

要素数を \(n\) として、構築・空間は \(O(n)\)、更新と和は \(O(\log n)\)、size は \(O(1)\) です。T の演算は \(O(1)\) とします。

検証コード

• test/lca_path_sum.test.cpp
• test/offline_point_set_range_frequency.test.cpp
• test/offline_range_frequency.test.cpp
• test/offline_range_sum.test.cpp
• test/offline_rectangle_sum.test.cpp
• test/weighted_wavelet_matrix_rectangle.test.cpp

コード

#ifndef FENWICK_SUM_HPP
#define FENWICK_SUM_HPP

#include <cassert>
#include <bit>
#include <limits>
#include <vector>

template <class T>
class FenwickSum {
	std::vector<T> data;

  public:
	FenwickSum() = default;
	explicit FenwickSum(int n) { assert(n >= 0); data.resize(n); }
	explicit FenwickSum(const std::vector<T> &a) { build(a); }
	void build(const std::vector<T> &a) {
		assert(a.size() <= static_cast<unsigned>(std::numeric_limits<int>::max()));
		data = a;
		constexpr unsigned step = sizeof(T) <= 8 ? 2 : 1;
		for (unsigned i = step; i <= data.size(); i += step) {
			if constexpr (step == 2) data[i - 1] += data[i - 2];
			unsigned p = i + (i & -i);
			if (p <= data.size()) data[p - 1] += data[i - 1];
		}
	}
	int size() const { return static_cast<int>(data.size()); }
	void add(int p, const T &x) {
		assert(0 <= p && p < size());
		for (unsigned i = unsigned(p) + 1; i <= data.size(); i += i & -i) data[i - 1] += x;
	}
	void add_pair(int p, const T &x, int q, const T &y) {
		assert(0 <= p && p < size() && 0 <= q && q < size());
		unsigned a = unsigned(p) + 1, b = unsigned(q) + 1;
		while (a != b && (a <= data.size() || b <= data.size())) {
			if (a < b) data[a - 1] += x, a += a & -a;
			else data[b - 1] += y, b += b & -b;
		}
		T both = x;
		both += y;
		for (; a <= data.size(); a += a & -a) data[a - 1] += both;
	}
	T prefix(int r) const {
		assert(0 <= r && r <= size());
		T result{};
		for (; r; r &= r - 1) result += data[r - 1];
		return result;
	}
	template <class Accum>
	Accum prod_as(int l, int r) const {
		assert(0 <= l && l <= r && r <= size());
		Accum result{};
		if (l == r) return result;
		unsigned shared = unsigned(l) & (std::numeric_limits<unsigned>::max() << std::bit_width(unsigned(l ^ r)));
		for (; unsigned(l) != shared; l &= l - 1) result -= data[l - 1];
		for (; unsigned(r) != shared; r &= r - 1) result += data[r - 1];
		return result;
	}
	T prod(int l, int r) const { return prod_as<T>(l, r); }
};

#endif

Back to top page