cp-library

競技プログラミング用 C++ ライブラリ

View the Project on GitHub qandori/cp-library

:heavy_check_mark: 永続セグメント木

data_structure/persistent_segtree.hpp · GitHub で見る

共通ページ

配列の過去の状態を残したまま、一点更新・区間積・境界探索を行います。任意の過去版から分岐でき、別の版の区間をコピーすることもできます。遅延伝搬による区間更新は扱いません。

#include "lib/data_structure/persistent_segtree.hpp"
#include <cassert>
#include <vector>

long long op(long long a, long long b) { return a + b; }
long long e() { return 0; }
using Tree = PersistentSegtree<long long, op, e>;

int main() {
	Tree seg(std::vector<long long>{1, 2, 3, 4});
	seg.reserve(1000);
	auto v0 = seg.initial_root();
	auto v1 = seg.set(v0, 1, 10);
	auto v2 = seg.multiply(v0, 3, 5);
	auto v3 = seg.copy_range(v1, v2, 2, 4);
	assert(seg.prod(v0, 0, 4) == 10);
	assert(seg.prod(v1, 0, 4) == 18);
	assert(seg.get(v2, 3) == 9);
	assert(seg.all_prod(v3) == 23);
	auto r = seg.max_right(v3, 0, [](long long s) { return s <= 14; });
	assert(r == 3);
}

v0 は最初の配列、v1 は位置 1 を 10 にした版、v2 は最初の配列の位置 3 に 5 を加えた版です。v1 と v2 は互いに独立しています。v3 は v1 の区間 [2,4) だけを v2 の値で置き換えた版です。更新を呼んでも元の版は変わらず、返されたルート ID を保存して新しい版を使います。

op(a,b) は結合的な演算、e() はその単位元です。和以外に最小値・最大値・行列積・関数合成などを扱えます。区間積は添字の小さい順に計算し、可換性は要求しません。S はコピー・ムーブ・代入できる型とし、デフォルト構築と等値比較は不要です。

操作 内容・返り値 計算量
PersistentSegtree(n = 0) 長さ n、全要素が e() の初期版を用意。実ノードはまだ作らない \(O(1)\)
PersistentSegtree(a) std::vector<S> から初期版を構築 \(O(n)\)
initial_root() コンストラクタで用意した初期版のルート ID \(O(1)\)
set(root,p,x) 版 root の位置 p を x に代入した新しいルート ID 償却 \(O(\log n)\)
multiply(root,p,x) 位置 p を op(現在値,x) にした新しいルート ID 償却 \(O(\log n)\)
get(root,p) 版 root の位置 p の値 \(O(\log n)\)
prod(root,l,r) 版 root の半開区間 [l,r) の積。空区間は e() \(O(\log n)\)
all_prod(root) 版 root の全体の積 \(O(1)\)
copy_range(root,source,l,r) root の [l,r) を source の同じ区間で置き換えた新しいルート ID 償却 \(O(\log n)\)
max_right(root,l,f) f(prod(root,l,r)) が真である最大の右端 r \(O(\log n)\)
min_left(root,r,f) f(prod(root,l,r)) が真である最小の左端 l \(O(\log n)\)
size() / empty() 配列の長さ/長さが 0 か \(O(1)\)
reserve(count) 空ノードを除くノード count 個分の容量を予約。既存のルート ID は変わらない 再確保時は現在のノード数に比例
node_count() 空ノードを除く、全版で作成したノード数 \(O(1)\)
memory_usage() 予約済みノード領域のバイト数。内部の空ノード 1 個分も含む \(O(1)\)

ルート ID の型は root_type(std::uint32_t)、値の型は value_type(S)です。ルート ID を変数や std::vector<Tree::root_type> にコピーするだけなら、時間・追加領域ともに一定です。返される ID は連番の「版番号」ではないので、版番号から ID への対応は利用者が保存してください。

添字は 0 始まりで、0 <= n <= INT_MAX、一点操作は 0 <= p < n、区間操作は 0 <= l <= r <= n とします。ルート ID はその木の initial_root・更新・区間コピーが返したものを使います。別の木の ID を混ぜることはできません。木のオブジェクト自体をコピーするとノード領域全体を複製します。移動元は再代入してから使ってください。

境界探索では f(e()) が真で、区間を伸ばすと真から偽へ単調に変わる必要があります。例えば非負の配列なら「和が指定値以下」を使えます。負の値がある配列の和では、この条件を満たすとは限りません。

get・prod・all_prod・境界探索は読み取り専用で、ノードを増やしません。区間コピーは部分木を共有し、全区間のコピー、空区間のコピー、同じ版からのコピーはノードを作りません。版ごとのノード解放は行わず、使わなくなった版のノードも木を破棄するまで保持します。

長さが正のとき、\(h=\lceil\log_4 n\rceil\) とすると、配列からの構築は \(\displaystyle\sum_{i=0}^{h}\lceil n/4^i\rceil\) ノード、一点更新は \(h+1\) ノードを作ります。reserve の引数は更新回数ではなく、既存分を含むノード数なので、更新回数からこの上限を見積もって予約できます。ノード数の上限は \(2^{32}-2\) 個です。

演算・値のコピーを \(O(1)\) とした計算量です。構築は \(O(n)\)、一点更新・区間積・境界探索・区間コピーは \(O(\log n)\)、更新の領域確保を含む時間は償却計算量です。長さ 0・1 の操作は \(O(1)\) です。構築後に更新・区間コピーを合計 \(u\) 回行うと、空間は \(O(n+u\log n)\)、長さだけを指定した疎な初期版では \(O(u\log n)\) です。予約容量が大きい場合はその容量分を使います。

検証コード

コード

#ifndef PERSISTENT_SEGTREE_HPP
#define PERSISTENT_SEGTREE_HPP

#include <array>
#include <bit>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <type_traits>
#include <utility>
#include <vector>

template <class S, auto op, auto e>
class PersistentSegtree {
  public:
	using value_type = S;
	using root_type = std::uint32_t;

  private:
	struct Node { S value; std::array<root_type, 4> child; };
	int n_ = 0, height = 0;
	root_type initial_ = 0;
	std::vector<Node> nodes{{e(), {}}};
	void check([[maybe_unused]] root_type root) const { assert(root < nodes.size()); }
	std::uint64_t capacity() const { return std::uint64_t(1) << (2 * height); }
	void initialize(int n) {
		assert(n >= 0);
		n_ = n;
		height = n ? (std::bit_width(static_cast<unsigned>(n - 1)) + 1) / 2 : 0;
	}
	root_type make(S value, std::array<root_type, 4> child = {}) {
		assert(nodes.size() < std::numeric_limits<root_type>::max());
		auto result = static_cast<root_type>(nodes.size());
		nodes.push_back({std::move(value), child});
		return result;
	}
	S combine(root_type a, root_type b) const {
		if (!a) return nodes[b].value;
		if (!b) return nodes[a].value;
		return op(nodes[a].value, nodes[b].value);
	}
	root_type join(std::array<root_type, 4> child) {
		if constexpr (std::is_trivially_copyable_v<S>) {
			return make(op(op(nodes[child[0]].value, nodes[child[1]].value), op(nodes[child[2]].value, nodes[child[3]].value)), child);
		} else {
			if (!(child[2] | child[3])) return make(combine(child[0], child[1]), child);
			if (!(child[0] | child[1])) return make(combine(child[2], child[3]), child);
			return make(op(combine(child[0], child[1]), combine(child[2], child[3])), child);
		}
	}
	root_type build(const std::vector<S> &a, std::uint64_t l, std::uint64_t width) {
		if (l >= a.size()) return 0;
		if (width == 1) return make(a[l]);
		width >>= 2;
		std::array<root_type, 4> child;
		for (int i = 0; i < 4; ++i) child[i] = build(a, l + i * width, width);
		return join(child);
	}
	template <bool Append>
	root_type update(root_type root, int p, S value) {
		check(root); assert(0 <= p && p < n_);
		std::array<std::array<root_type, 4>, 16> path;
		for (int level = height - 1; level >= 0; --level) {
			path[level] = nodes[root].child;
			root = path[level][(static_cast<unsigned>(p) >> (2 * level)) & 3];
		}
		if constexpr (Append) value = op(nodes[root].value, value);
		root = make(std::move(value));
		for (int level = 0; level < height; ++level) {
			path[level][(static_cast<unsigned>(p) >> (2 * level)) & 3] = root;
			root = join(path[level]);
		}
		return root;
	}
	S prefix(root_type root, int level, std::uint64_t end) const {
		auto width = std::uint64_t(1) << (2 * level);
		S result = e();
		while (root && end < width) {
			width >>= 2;
			unsigned d = static_cast<unsigned>(end >> (2 * --level));
			const auto &child = nodes[root].child;
			if (d >= 2) result = op(std::move(result), op(nodes[child[0]].value, nodes[child[1]].value));
			if (d & 1) result = op(std::move(result), nodes[child[d - 1]].value);
			end -= d * width;
			if (!end) return result;
			root = child[d];
		}
		return root ? op(std::move(result), nodes[root].value) : result;
	}
	S suffix(root_type root, int level, std::uint64_t start) const {
		auto width = std::uint64_t(1) << (2 * level);
		S result = e();
		while (root && start) {
			width >>= 2;
			unsigned d = static_cast<unsigned>(start >> (2 * --level));
			const auto &child = nodes[root].child;
			if (d < 2) result = op(op(nodes[child[2]].value, nodes[child[3]].value), std::move(result));
			if (!(d & 1)) result = op(nodes[child[d + 1]].value, std::move(result));
			root = child[d];
			start -= d * width;
		}
		return root ? op(nodes[root].value, std::move(result)) : result;
	}
	root_type copy_impl(root_type root, root_type source, std::uint64_t l, std::uint64_t width, std::uint64_t ql, std::uint64_t qr) {
		if (root == source) return root;
		if (ql <= l && l + width <= qr) return source;
		width >>= 2;
		auto old = nodes[root].child, from = nodes[source].child, child = old;
		for (int i = 0; i < 4; ++i) {
			auto start = l + i * width;
			if (start < qr && ql < start + width) child[i] = copy_impl(old[i], from[i], start, width, ql, qr);
		}
		if (child == old) return root;
		if (child == from) return source;
		return join(child);
	}
	template <bool Right, class Pred>
	std::uint64_t search(root_type root, std::uint64_t l, std::uint64_t width, std::uint64_t end, S &value, Pred &f) const {
		if (!root) return Right ? capacity() : 0;
		if (Right ? end <= l : l + width <= end) {
			S next = Right ? op(value, nodes[root].value) : op(nodes[root].value, value);
			if (f(next)) { value = std::move(next); return Right ? capacity() : 0; }
			if (width == 1) return Right ? l : l + 1;
		}
		width >>= 2;
		for (int step = 0; step < 4; ++step) {
			int i = Right ? step : 3 - step;
			auto start = l + i * width;
			if (Right ? start + width <= end : end <= start) continue;
			auto found = search<Right>(nodes[root].child[i], start, width, end, value, f);
			if (Right ? found != capacity() : found != 0) return found;
		}
		return Right ? capacity() : 0;
	}

  public:
	explicit PersistentSegtree(int n = 0) { initialize(n); }
	explicit PersistentSegtree(const std::vector<S> &a) {
		assert(a.size() <= static_cast<std::size_t>(std::numeric_limits<int>::max()));
		initialize(static_cast<int>(a.size()));
		if (n_) {
			std::size_t count = 0;
			for (auto width = capacity(); width; width >>= 2) count += (a.size() + width - 1) / width;
			reserve(count);
			initial_ = build(a, 0, capacity());
		}
	}
	int size() const { return n_; }
	bool empty() const { return !n_; }
	root_type initial_root() const { return initial_; }
	void reserve(std::size_t count) { assert(count < std::numeric_limits<root_type>::max()); nodes.reserve(count + 1); }
	std::size_t node_count() const { return nodes.size() - 1; }
	std::size_t memory_usage() const { return nodes.capacity() * sizeof(Node); }
	root_type set(root_type root, int p, S value) { return update<false>(root, p, std::move(value)); }
	root_type multiply(root_type root, int p, S value) { return update<true>(root, p, std::move(value)); }
	S get(root_type root, int p) const {
		check(root); assert(0 <= p && p < n_);
		for (int level = height - 1; root && level >= 0; --level) root = nodes[root].child[(static_cast<unsigned>(p) >> (2 * level)) & 3];
		return nodes[root].value;
	}
	S prod(root_type root, int l, int r) const {
		check(root); assert(0 <= l && l <= r && r <= n_);
		if (l == r) return e();
		if (!l && r == n_) return nodes[root].value;
		auto first = static_cast<std::uint64_t>(l), last = static_cast<std::uint64_t>(r), width = capacity();
		int level = height;
		while (root) {
			if (!first && last == width) return nodes[root].value;
			width >>= 2;
			int shift = 2 * --level;
			unsigned a = static_cast<unsigned>(first >> shift), b = static_cast<unsigned>((last - 1) >> shift);
			if (a == b) { root = nodes[root].child[a]; first -= a * width; last -= a * width; continue; }
			const auto &child = nodes[root].child;
			S result = suffix(child[a], level, first - a * width);
			for (unsigned i = a + 1; i < b; ++i) result = op(std::move(result), nodes[child[i]].value);
			return op(std::move(result), prefix(child[b], level, last - b * width));
		}
		return e();
	}
	S all_prod(root_type root) const { check(root); return nodes[root].value; }
	root_type copy_range(root_type root, root_type source, int l, int r) {
		check(root); check(source); assert(0 <= l && l <= r && r <= n_);
		if (l == r || root == source) return root;
		if (!l && r == n_) return source;
		return copy_impl(root, source, 0, capacity(), l, r);
	}
	template <class Pred> int max_right(root_type root, int l, Pred f) const {
		check(root); assert(0 <= l && l <= n_ && f(e()));
		S value = e();
		auto result = l == n_ ? static_cast<std::uint64_t>(n_) : search<true>(root, 0, capacity(), l, value, f);
		return result < static_cast<unsigned>(n_) ? static_cast<int>(result) : n_;
	}
	template <class Pred> int min_left(root_type root, int r, Pred f) const {
		check(root); assert(0 <= r && r <= n_ && f(e()));
		S value = e();
		return r ? static_cast<int>(search<false>(root, 0, capacity(), r, value, f)) : 0;
	}
};

#endif
Back to top page