Segment tree with lazy propagation (src/datastructure/lazy_segment_tree.hpp)

• View this file on GitHub
• View document part on GitHub
• Last update: 2024-01-06 12:34:09+09:00
• Include: #include "src/datastructure/lazy_segment_tree.hpp"

概要

区間更新・区間取得を $O(\log N)$ で行えます。

要求される性質

遅延伝播セグメント木で区間取得の際に扱う要素と、それに対して作用を与える区間更新の際に扱う要素は、どちらもモノイドである必要があります¹。

区間取得の際に扱うモノイドを $(S, \cdot, e_S)$ とし、それに対して作用を与える区間更新の際に扱うモノイドを $(T, \circ, e_T)$ とします。そして $T$ の $S$ への作用を外部二項演算 $\ast: T \times S \rightarrow S$ とします。 このとき、外部二項演算 $\ast$ は準同型であり、モノイド作用である必要があります。 つまり、次の条件を満たす必要があります。

• モノイド作用
  • $e_T \ast s = s \quad (\forall s \in S)$ を満たす
  • $t_1 \ast (t_2 \ast s) = (t_1 \circ t_2) \ast s \quad (\forall s \in S, \forall t_1, t_2 \in T)$ を満たす
• 準同型
  • $t \ast (s_1 \cdot s_2) = (t \ast s_1) \cdot (t \ast s_2) \quad (\forall s_1, s_2 \in S, \forall t \in T)$ を満たす²

1. 実際にはモノイドよりも条件を緩めることができます。 ↩

2. $t \ast s = f_t(s)$ とすると $f_{t}(s_1 \cdot s_2) = f_{t}(s_1) \cdot f_{t}(s_2)$ となるので、準同型を表していることが分かりやすくなります。 ↩

Verified with

• test/datastructure/lazy_segment_tree/aoj_dsl_2_d.test.cpp
• test/datastructure/lazy_segment_tree/aoj_dsl_2_f.test.cpp

Code

#ifndef SEGTREE_LAZY_SEGMENT_TREE_HPP
#define SEGTREE_LAZY_SEGMENT_TREE_HPP

#include <vector>

template<class T, class E, class F, class G, class H>
class LazySegmentTree { // 0-indexed
private:
	int n_{}, height{};
	std::vector<T> tree;
	std::vector<E> lazy;
	F f; // function<T(T, T)>
	G g; // function<T(T, E)>
	H h; // function<E(E, E)>
	T ti;
	E ei;

	inline T reflect(int k) { return (lazy[k] == ei ? tree[k] : g(tree[k], lazy[k])); }

	inline void eval(int k) {
		if (lazy[k] == ei) return;
		lazy[2 * k] = h(lazy[2 * k], lazy[k]);
		lazy[2 * k + 1] = h(lazy[2 * k + 1], lazy[k]);
		tree[k] = reflect(k);
		lazy[k] = ei;
	}

	inline void thrust(int k) {
		for (int i = height; i > 0; --i) eval(k >> i);
	}

	inline void recalc(int k) {
		while (k >>= 1) tree[k] = f(reflect(2 * k), reflect(2 * k + 1));
	}

public:
	LazySegmentTree(F f, G g, H h, T ti, E ei) : f(f), g(g), h(h), ti(ti), ei(ei) {}

	void init(int n) {
		n_ = 1, height = 0;
		while (n_ < n) n_ *= 2, ++height;
		tree.assign(2 * n_, ti);
		lazy.assign(2 * n_, ei);
	}

	void build(std::vector<T> const &v) {
		int const N = v.size();
		init(N);
		for (int i = 0; i < N; ++i) tree[n_ + i] = v[i];
		for (int i = n_ - 1; i > 0; --i) tree[i] = f(tree[2 * i], tree[2 * i + 1]);
	}

	void update(int s, int t, E const &x) { // [l, r)
		s += n_, t += n_;
		thrust(s), thrust(t - 1);
		int l = s, r = t;
		while (l < r) {
			if (l & 1) lazy[l] = h(lazy[l], x), ++l;
			if (r & 1) --r, lazy[r] = h(lazy[r], x);
			l >>= 1, r >>= 1;
		}
		recalc(s), recalc(t - 1);
	}

	T find(int s, int t) { // [l, r)
		s += n_, t += n_;
		thrust(s), thrust(t - 1);
		int l = s, r = t;
		T ll = ti, rr = ti;
		while (l < r) {
			if (l & 1) ll = f(ll, reflect(l++));
			if (r & 1) rr = f(reflect(--r), rr);
			l >>= 1, r >>= 1;
		}
		return f(ll, rr);
	}

	T at(int i) {
		i += n_;
		thrust(i);
		return reflect(i);
	}
};

#endif // SEGTREE_LAZY_SEGMENT_TREE_HPP

#line 1 "src/datastructure/lazy_segment_tree.hpp"



#include <vector>

template<class T, class E, class F, class G, class H>
class LazySegmentTree { // 0-indexed
private:
	int n_{}, height{};
	std::vector<T> tree;
	std::vector<E> lazy;
	F f; // function<T(T, T)>
	G g; // function<T(T, E)>
	H h; // function<E(E, E)>
	T ti;
	E ei;

	inline T reflect(int k) { return (lazy[k] == ei ? tree[k] : g(tree[k], lazy[k])); }

	inline void eval(int k) {
		if (lazy[k] == ei) return;
		lazy[2 * k] = h(lazy[2 * k], lazy[k]);
		lazy[2 * k + 1] = h(lazy[2 * k + 1], lazy[k]);
		tree[k] = reflect(k);
		lazy[k] = ei;
	}

	inline void thrust(int k) {
		for (int i = height; i > 0; --i) eval(k >> i);
	}

	inline void recalc(int k) {
		while (k >>= 1) tree[k] = f(reflect(2 * k), reflect(2 * k + 1));
	}

public:
	LazySegmentTree(F f, G g, H h, T ti, E ei) : f(f), g(g), h(h), ti(ti), ei(ei) {}

	void init(int n) {
		n_ = 1, height = 0;
		while (n_ < n) n_ *= 2, ++height;
		tree.assign(2 * n_, ti);
		lazy.assign(2 * n_, ei);
	}

	void build(std::vector<T> const &v) {
		int const N = v.size();
		init(N);
		for (int i = 0; i < N; ++i) tree[n_ + i] = v[i];
		for (int i = n_ - 1; i > 0; --i) tree[i] = f(tree[2 * i], tree[2 * i + 1]);
	}

	void update(int s, int t, E const &x) { // [l, r)
		s += n_, t += n_;
		thrust(s), thrust(t - 1);
		int l = s, r = t;
		while (l < r) {
			if (l & 1) lazy[l] = h(lazy[l], x), ++l;
			if (r & 1) --r, lazy[r] = h(lazy[r], x);
			l >>= 1, r >>= 1;
		}
		recalc(s), recalc(t - 1);
	}

	T find(int s, int t) { // [l, r)
		s += n_, t += n_;
		thrust(s), thrust(t - 1);
		int l = s, r = t;
		T ll = ti, rr = ti;
		while (l < r) {
			if (l & 1) ll = f(ll, reflect(l++));
			if (r & 1) rr = f(reflect(--r), rr);
			l >>= 1, r >>= 1;
		}
		return f(ll, rr);
	}

	T at(int i) {
		i += n_;
		thrust(i);
		return reflect(i);
	}
};

Back to top page