セグメント木

data_structure/segtree.hpp · GitHub で見る

共通ページ

一点代入と区間集約を扱います。op は結合則を満たす演算、e() は単位元です。可換性は不要です。

#include "lib/data_structure/segtree.hpp"

long long op(long long a, long long b) { return a + b; }
long long e() { return 0; }

Segtree<long long, op, e> tree(std::vector<long long>{1, 2, 3});
tree.set(1, 5);
auto sum = tree.prod(0, 3); // 9
int r = tree.max_right(0, [](long long x) { return x <= 6; }); // 2

Segtree(n) は全要素を e() で初期化します。添字は 0 始まり、区間は半開区間です。

操作	内容・返り値	計算量
Segtree(n = 0) / Segtree(values)	長さ／値の配列から構築	\(O(n)\)
size()	要素数	\(O(1)\)
set(p,value) / get(p)	一点代入／取得	\(O(\log n)\) / \(O(1)\)
prod(l,r)	左から右の順で区間を集約。空区間なら e()	\(O(\log n)\)
all_prod()	全体の集約。空列なら e()	\(O(1)\)
max_right(l,pred)	[l,r) の集約が述語を満たす最大の r	\(O(\log n)\)
min_left(r,pred)	[l,r) の集約が述語を満たす最小の l	\(O(\log n)\)

境界検索は pred(e()) == true で、区間を伸ばすと一度だけ true から false に変わる述語を指定します。

長さは非負の int、一点引数は [0,n)、区間は 0 <= l <= r <= n が必要です。空列に対応し、コピーは独立します。更新以外の操作は const な木にも使えます。演算結果は指定した型で表現できるようにしてください。

長さを \(n\) として、構築・空間は \(O(n)\)、一点更新・区間集約・境界検索は \(O(\log n)\)、一点取得・全体集約・要素数の取得は \(O(1)\) です。演算・コピーは \(O(1)\) とします。

検証コード

• test/segtree_composite.test.cpp

コード

#ifndef SEGTREE_HPP
#define SEGTREE_HPP

#include <algorithm>
#include <bit>
#include <cassert>
#include <cstddef>
#include <limits>
#include <utility>
#include <vector>

template<class S, auto op, auto e> class Segtree {
	int n_ = 0;
	unsigned base_ = 1;
	std::vector<S> data_;
	void pull(unsigned k) { data_[k] = op(data_[k * 2], data_[k * 2 + 1]); }
	void allocate(int n) {
		assert(n >= 0);
		n_ = n;
		base_ = std::bit_ceil(unsigned(n) ? unsigned(n) : 1U);
		data_.assign(std::size_t(base_) * 2, e());
	}

  public:
	explicit Segtree(int n = 0) { allocate(n); }
	explicit Segtree(const std::vector<S> &a) {
		assert(a.size() <= std::size_t(std::numeric_limits<int>::max()));
		allocate(static_cast<int>(a.size()));
		for (int i = 0; i < n_; ++i) data_[base_ + unsigned(i)] = a[i];
		for (unsigned k = base_ - 1; k; --k) pull(k);
	}
	Segtree(const Segtree &) = default;
	Segtree &operator=(const Segtree &other) {
		if (this != &other) { Segtree copy(other); swap(copy); }
		return *this;
	}
	Segtree(Segtree &&other) noexcept
		: n_(std::exchange(other.n_, 0)), base_(std::exchange(other.base_, 1)), data_(std::move(other.data_)) {}
	Segtree &operator=(Segtree &&other) noexcept {
		if (this != &other) {
			n_ = std::exchange(other.n_, 0);
			base_ = std::exchange(other.base_, 1);
			data_ = std::move(other.data_);
		}
		return *this;
	}
	void swap(Segtree &other) noexcept {
		std::swap(n_, other.n_), std::swap(base_, other.base_);
		data_.swap(other.data_);
	}
	[[nodiscard]] int size() const { return n_; }
	void set(int p, S value) {
		assert(0 <= p && p < n_);
		unsigned k = base_ + unsigned(p);
		data_[k] = std::move(value);
		for (k >>= 1; k; k >>= 1) pull(k);
	}
	[[nodiscard]] S get(int p) const {
		assert(0 <= p && p < n_);
		return data_[base_ + unsigned(p)];
	}
	[[nodiscard]] S all_prod() const { return n_ ? data_[1] : e(); }
	[[nodiscard]] S prod(int l, int r) const {
		assert(0 <= l && l <= r && r <= n_);
		if (l == r) return e();
		unsigned a = base_ + unsigned(l), b = base_ + unsigned(r) - 1;
		unsigned level = std::bit_width(a ^ b);
		if (!level) return data_[a];
		unsigned left_stop = a >> (level - 1), right_stop = left_stop + 1;
		a >>= (std::min)(unsigned(std::countr_zero(a)), level - 1);
		b >>= (std::min)(unsigned(std::countr_one(b)), level - 1);
		S left = data_[a], right = data_[b];
		for (;;) {
			++a;
			a >>= std::countr_zero(a);
			if (a <= right_stop) break;
			left = op(std::move(left), data_[a]);
		}
		for (;;) {
			--b;
			b >>= std::countr_one(b);
			if (b <= right_stop) break;
			right = op(data_[b], std::move(right));
		}
		return op(std::move(left), std::move(right));
	}
	template<class Pred> [[nodiscard]] int max_right(int l, Pred f) const {
		assert(0 <= l && l <= n_ && f(e()));
		if (l == n_) return n_;
		unsigned k = base_ + unsigned(l);
		S value = e();
		do {
			while (!(k & 1)) k >>= 1;
			S next = op(value, data_[k]);
			if (!f(next)) {
				while (k < base_) {
					k *= 2;
					next = op(value, data_[k]);
					if (f(next)) value = std::move(next), ++k;
				}
				return static_cast<int>(k - base_);
			}
			value = std::move(next);
			++k;
		} while ((k & -k) != k);
		return n_;
	}
	template<class Pred> [[nodiscard]] int min_left(int r, Pred f) const {
		assert(0 <= r && r <= n_ && f(e()));
		if (!r) return 0;
		unsigned k = base_ + unsigned(r);
		S value = e();
		do {
			--k;
			while (k > 1 && (k & 1)) k >>= 1;
			S next = op(data_[k], value);
			if (!f(next)) {
				while (k < base_) {
					k = k * 2 + 1;
					next = op(data_[k], value);
					if (f(next)) value = std::move(next), --k;
				}
				return static_cast<int>(k + 1 - base_);
			}
			value = std::move(next);
		} while ((k & -k) != k);
		return 0;
	}
};

#endif

Back to top page