Fenwick Tree

data_structure/fenwick_tree.hpp · GitHub で見る

共通ページ

#include "lib/data_structure/fenwick_tree.hpp"

FenwickTree<long long> tree(std::vector<long long>{1, 2, 3, 4});
tree.add(1, 5);
auto sum = tree.sum(1, 3); // 10
auto prefix = tree.sum(3); // 11

FenwickTree<T>(n) は長さ n の零列、ベクトルを渡す構築はその列で初期化します。添字は 0 始まり、add(p,x) は一点加算、sum(l,r) は半開区間の和、sum(r) は [0,r) の和、size() は要素数です。空列・空区間にも対応します。コピーは独立します。

T は零初期化・加減算ができる可換群の型です。整数、modint などを指定できます。整数型は内部を符号なし整数にし、符号付きオーバーフローを避けます。結果は型のビット幅に応じた剰余です。bool は指定できません。浮動小数点型は和の順序によって丸め誤差が変わります。

長さは非負の int、更新位置は [0,n)、区間は 0 <= l <= r <= n が必要です。

両端の prefix に共通する祖先を読み飛ばします。既定では 1024 要素ごとに 1 要素分のパディングを入れ、キャッシュの競合を減らします。FenwickTree<T, false> で無効にできます。効果は区間長・CPU・コンパイル条件に依存します。

長さを \(n\) として、構築・空間は \(O(n)\)、加算・和は \(O(\log n)\)、size() は \(O(1)\) です。型の演算・コピーは \(O(1)\) とします。

検証コード

• test/fenwick_tree.test.cpp
• test/hld_path_sum.test.cpp
• test/hld_subtree_sum.test.cpp

コード

#ifndef FENWICK_TREE_HPP
#define FENWICK_TREE_HPP

#include <cassert>
#include <cstddef>
#include <limits>
#include <type_traits>
#include <vector>

namespace fenwick_tree_detail {
template<class T, bool = std::is_integral_v<T>> struct Storage { using type = T; };
template<class T> struct Storage<T, true> { using type = std::make_unsigned_t<T>; };
#if defined(__SIZEOF_INT128__)
template<> struct Storage<__int128_t, false> { using type = __uint128_t; };
#endif
}

template<class T, bool Padded = true> class FenwickTree {
	static_assert(!std::is_same_v<T, bool>);
	using U = typename fenwick_tree_detail::Storage<T>::type;
	std::vector<U> data_;
	int n_ = 0;
	static std::size_t index(unsigned i) {
		if constexpr (Padded) return std::size_t(i) + (i >> 10);
		else return i;
	}

  public:
	explicit FenwickTree(int n = 0) : n_(n) {
		assert(n >= 0);
		data_.resize(index(unsigned(n)) + 1);
	}
	explicit FenwickTree(const std::vector<T> &a) {
		assert(a.size() <= std::size_t(std::numeric_limits<int>::max()));
		n_ = static_cast<int>(a.size());
		data_.resize(index(unsigned(n_)) + 1);
		for (unsigned i = 2; i <= unsigned(n_); i += 2) {
			data_[index(i - 1)] += U(a[i - 2]);
			data_[index(i)] += data_[index(i - 1)];
			data_[index(i)] += U(a[i - 1]);
			unsigned p = i + (i & -i);
			if (p <= unsigned(n_)) data_[index(p)] += data_[index(i)];
		}
		if (n_ & 1) data_[index(unsigned(n_))] += U(a.back());
	}
	[[nodiscard]] int size() const { return n_; }
	void add(int p, T x) {
		assert(0 <= p && p < n_);
		for (unsigned i = unsigned(p) + 1; i <= unsigned(n_); i += i & -i) data_[index(i)] += U(x);
	}
	[[nodiscard]] T sum(int r) const {
		assert(0 <= r && r <= n_);
		U result{};
		for (unsigned i = unsigned(r); i; i &= i - 1) result += data_[index(i)];
		return T(result);
	}
	[[nodiscard]] T sum(int l, int r) const {
		assert(0 <= l && l <= r && r <= n_);
		U result{};
		unsigned a = unsigned(l), b = unsigned(r);
		while (b > a) result += data_[index(b)], b &= b - 1;
		while (a > b) result -= data_[index(a)], a &= a - 1;
		return T(result);
	}
};

#endif

Back to top page