Use lazy segment tree from el vasito (#6)

* Rename `LazySegmentTree` to `MemoryLazySegmentTree`
* Add new `LazySegmentTree`

content/data-structures/LazySegmentTree.h

@@ -1,69 +1,70 @@
 /**
- * Author: Simon Lindholm
- * Date: 2016-10-08
+ * Author: Iván Renison
+ * Date: 2024-03-05
  * License: CC0
- * Source: me
- * Description: Segment tree with ability to add or set values of large intervals, and compute max of intervals.
+ * Source: notebook el vasito
+ * Description: Segment tree with ability to add values of large intervals, and compute the sum of intervals.
+ * Ranges are [s, e).
  * Can be changed to other things.
- * Use with a bump allocator for better performance, and SmallPtr or implicit indices to save memory.
  * Time: O(\log N).
- * Usage: Node* tr = new Node(v, 0, sz(v));
- * Status: stress-tested a bit
+ * Usage: STree st(n);
+ *  st.init(x);
+ *  st.upd(s, e, v);
+ *  st.query(s, e);
+ * Status: Tested on SPOJ HORRIBLE, stress-tested a bit
  */
 #pragma once
 
-#include "../various/BumpAllocator.h"
+struct STree { // example: range sum with range addition
+	typedef ll T; typedef ll L; // T: data type, L: lazy type
+	constexpr static T tneut = 0; constexpr static L lneut = 0;
+	// neutrals
+	T f(T a, T b) { return a + b; } // operation
+	T apply(T v, L l, ll s, ll e) { return v + l * (e - s); }
+	// new st according to lazy
+	L comb(L a, L b) { return a + b; }
+	// cumulative effect of lazy
 
-const ll inf = 1e18;
-struct Node {
-	typedef ll T; // data type
-	struct L { ll mset, madd; }; // lazy type
-	const T tneut = -inf;     // neutral elements
-	const L lneut = {inf, 0};
-	T f (T a, T b) { return max(a, b); } // (any associative fn)
-	T apply (T a, L b) {
-		return b.mset != inf ? b.mset + b.madd : a + b.madd;
-	} // Apply lazy
-	L comb(L a, L b) {
-		if (b.mset != inf) return b;
-		return {a.mset, a.madd + b.madd};
-	} // Combine lazy
-
-	Node *l = 0, *r = 0;
-	ll lo, hi; T val = tneut; L lazy = lneut;
-	Node(ll lo,ll hi):lo(lo),hi(hi){}//Large interval of tneut
-	Node(vector<T>& v, ll lo, ll hi) : lo(lo), hi(hi) {
-		if (lo + 1 < hi) {
-			ll mid = lo + (hi - lo)/2;
-			l = new Node(v, lo, mid); r = new Node(v, mid, hi);
-			val = f(l->val, r->val);
-		}
-		else val = v[lo];
+	vector<T> st;
+	vector<L> lazy;
+	ll n;
+	STree(ll n) : st(4*n, tneut), lazy(4*n, lneut), n(n) {}
+	void init(ll k, ll s, ll e, const vector<T> &a) {
+		lazy[k] = lneut;
+		if (s + 1 == e) { st[k] = a[s]; return; }
+		ll m = (s + e) / 2;
+		init(2*k, s, m, a), init(2*k+1, m, e, a);
+		st[k] = f(st[2*k], st[2*k+1]);
 	}
-	T query(ll L, ll R) {
-		if (R <= lo || hi <= L) return tneut;
-		if (L <= lo && hi <= R) return apply(val, lazy);
-		push();
-		return f(l->query(L, R), r->query(L, R));
-	}
-	void upd(ll Le, ll Ri, L x) {
-		if (Ri <= lo || hi <= Le) return;
-		if (Le <= lo && hi <= Ri) lazy = comb(lazy, x);
-		else {
-			push(), l->upd(Le, Ri, x), r->upd(Le, Ri, x);
-			val = f(l->query(lo, hi), r->query(lo, hi));
+	void push(ll k, ll s, ll e) {
+		if (lazy[k] == lneut) return; // if neutral, nothing to do
+		st[k] = apply(st[k], lazy[k], s, e);
+		if (s + 1 < e) { // propagate to children
+			lazy[2*k] = comb(lazy[2*k], lazy[k]);
+			lazy[2*k+1] = comb(lazy[2*k+1], lazy[k]);
 		}
+		lazy[k] = lneut; // clear node lazy
 	}
-	void set(ll L, ll R, ll x) { upd(L, R, {x, 0}); }
-	void add(ll L, ll R, ll x) { upd(L, R, {inf, x}); }
-	void push() {
-		if (!l) {
-			ll mid = lo + (hi - lo)/2;
-			l = new Node(lo, mid), r = new Node(mid, hi);
+	void upd(ll k, ll s, ll e, ll a, ll b, L v) {
+		push(k, s, e);
+		if (s >= b || e <= a) return;
+		if (s >= a && e <= b) {
+			lazy[k] = comb(lazy[k], v); // accumulate lazy
+			push(k, s, e);
+			return;
 		}
-		l->lazy = comb(l->lazy, lazy);
-		r->lazy = comb(r->lazy, lazy);
-		lazy = lneut;
-		val = f(l->query(lo, hi), r->query(lo, hi));
+		ll m = (s + e) / 2;
+		upd(2*k, s, m, a, b, v), upd(2*k+1, m, e, a, b, v);
+		st[k] = f(st[2*k], st[2*k+1]);
+	}
+	T query(ll k, ll s, ll e, ll a, ll b) {
+		if (s >= b || e <= a) return tneut;
+		push(k, s, e);
+		if (s >= a && e <= b) return st[k];
+		ll m = (s + e) / 2;
+		return f(query(2*k, s, m, a, b),query(2*k+1, m, e, a, b));
 	}
+	void init(const vector<T> &a) { init(1, 0, n, a); }
+	void upd(ll a, ll b, L v) { upd(1, 0, n, a, b, v); }
+	T query(ll a, ll b) { return query(1, 0, n, a, b); }
 };

content/data-structures/MemoryLazySegmentTree.h

@@ -0,0 +1,69 @@
+/**
+ * Author: Simon Lindholm
+ * Date: 2016-10-08
+ * License: CC0
+ * Source: me
+ * Description: Segment tree with ability to add or set values of large intervals, and compute max of intervals.
+ * Can be changed to other things.
+ * Use with a bump allocator for better performance, and SmallPtr or implicit indices to save memory.
+ * Time: O(\log N).
+ * Usage: Node* tr = new Node(v, 0, sz(v));
+ * Status: stress-tested a bit
+ */
+#pragma once
+
+#include "../various/BumpAllocator.h"
+
+const ll inf = 1e18;
+struct Node {
+	typedef ll T; // data type
+	struct L { ll mset, madd; }; // lazy type
+	const T tneut = -inf;     // neutral elements
+	const L lneut = {inf, 0};
+	T f (T a, T b) { return max(a, b); } // (any associative fn)
+	T apply (T a, L b) {
+		return b.mset != inf ? b.mset + b.madd : a + b.madd;
+	} // Apply lazy
+	L comb(L a, L b) {
+		if (b.mset != inf) return b;
+		return {a.mset, a.madd + b.madd};
+	} // Combine lazy
+
+	Node *l = 0, *r = 0;
+	ll lo, hi; T val = tneut; L lazy = lneut;
+	Node(ll lo,ll hi):lo(lo),hi(hi){}//Large interval of tneut
+	Node(vector<T>& v, ll lo, ll hi) : lo(lo), hi(hi) {
+		if (lo + 1 < hi) {
+			ll mid = lo + (hi - lo)/2;
+			l = new Node(v, lo, mid); r = new Node(v, mid, hi);
+			val = f(l->val, r->val);
+		}
+		else val = v[lo];
+	}
+	T query(ll L, ll R) {
+		if (R <= lo || hi <= L) return tneut;
+		if (L <= lo && hi <= R) return apply(val, lazy);
+		push();
+		return f(l->query(L, R), r->query(L, R));
+	}
+	void upd(ll Le, ll Ri, L x) {
+		if (Ri <= lo || hi <= Le) return;
+		if (Le <= lo && hi <= Ri) lazy = comb(lazy, x);
+		else {
+			push(), l->upd(Le, Ri, x), r->upd(Le, Ri, x);
+			val = f(l->query(lo, hi), r->query(lo, hi));
+		}
+	}
+	void set(ll L, ll R, ll x) { upd(L, R, {x, 0}); }
+	void add(ll L, ll R, ll x) { upd(L, R, {inf, x}); }
+	void push() {
+		if (!l) {
+			ll mid = lo + (hi - lo)/2;
+			l = new Node(lo, mid), r = new Node(mid, hi);
+		}
+		l->lazy = comb(l->lazy, lazy);
+		r->lazy = comb(r->lazy, lazy);
+		lazy = lneut;
+		val = f(l->query(lo, hi), r->query(lo, hi));
+	}
+};

content/data-structures/chapter.tex

@@ -4,6 +4,7 @@ \chapter{Data structures}
 \kactlimport{HashMap.h}
 \kactlimport{SegmentTree.h}
 \kactlimport{LazySegmentTree.h}
+\kactlimport{MemoryLazySegmentTree.h}
 \kactlimport{UnionFind.h}
 \kactlimport{UnionFindRollback.h}
 \kactlimport{SubMatrix.h}

content/graph/HLD.h

@@ -15,7 +15,7 @@
  */
 #pragma once
 
-#include "../data-structures/LazySegmentTree.h"
+#include "../data-structures/MemoryLazySegmentTree.h"
 
 template <bool VALS_EDGES> struct HLD {
 	ll N, tim = 0;

stress-tests/data-structures/LazySegmentTree.cpp

@@ -2,45 +2,60 @@
 
 #include "../../content/data-structures/LazySegmentTree.h"
 
-static unsigned R;
-ll ra() {
-	R *= 791231;
-	R += 1231;
-	return (ll)(R >> 1);
-}
+// Slow version with the same interface
+struct Slow {
+	vi elems;
+	ll n;
+	Slow(ll n) : elems(n), n(n) {}
+	void init(const vector<ll> &a) {
+		elems = a;
+	}
+	void upd(ll a, ll b, ll v) {
+		fore(i, a, b) {
+			elems[i] += v;
+		}
+	}
+	ll query(ll a, ll b) {
+		ll ans = 0;
+		fore(i, a, b) {
+			ans += elems[i];
+		}
+		return ans;
+	}
+};
 
-volatile ll res;
 int main() {
-	ll N = 10;
-	vi v(N);
-	iota(all(v), 0);
-	random_shuffle(all(v), [](ll x) { return ra() % x; });
-	Node* tr = new Node(v,0,N);
-	rep(i,0,N) rep(j,0,N) if (i <= j) {
-		ll ma = -inf;
-		rep(k,i,j) ma = max(ma, v[k]);
-		assert(ma == tr->query(i,j));
-	}
-	rep(it,0,1000000) {
-		ll i = ra() % (N+1), j = ra() % (N+1);
-		if (i > j) swap(i, j);
-		ll x = (ra() % 10) - 5;
 
-		ll r = ra() % 100;
-		if (r < 30) {
-			::res = tr->query(i, j);
-			ll ma = -inf;
-			rep(k,i,j) ma = max(ma, v[k]);
-			assert(ma == ::res);
-		}
-		else if (r < 70) {
-			tr->add(i, j, x);
-			rep(k,i,j) v[k] += x;
+	fore(_, 0, 100) {
+		ll n = rand() % 99 + 1;
+		vector<ll> a(n);
+		fore(i, 0, n) {
+			a[i] = rand() * (rand() % 2 ? 1 : -1);
 		}
-		else {
-			tr->set(i, j, x);
-			rep(k,i,j) v[k] = x;
+		STree st(n);
+		assert(st.n == n);
+		assert(st.query(0, n) == 0);
+		Slow slow(n);
+		st.init(a);
+		slow.init(a);
+
+		assert(st.query(0, n) == slow.query(0, n));
+
+		fore(_, 0, 100) {
+			ll t = rand() % 2;
+			ll l = rand() % n, r = rand() % n;
+			if (l > r) swap(l, r);
+			if (t == 0) {
+				ll v = rand() * (rand() % 2 ? 1 : -1);
+				st.upd(l, r, v);
+				slow.upd(l, r, v);
+			} else {
+				ll ans = st.query(l, r);
+				ll slowAns = slow.query(l, r);
+				assert(ans == slowAns);
+			}
 		}
 	}
-	cout<<"Tests passed!"<<endl;
+
+	cout << "Tests passed!" << endl;
 }

stress-tests/data-structures/MemoryLazySegmentTree.cpp

@@ -0,0 +1,46 @@
+#include "../utilities/template.h"
+
+#include "../../content/data-structures/MemoryLazySegmentTree.h"
+
+static unsigned R;
+ll ra() {
+	R *= 791231;
+	R += 1231;
+	return (ll)(R >> 1);
+}
+
+volatile ll res;
+int main() {
+	ll N = 10;
+	vi v(N);
+	iota(all(v), 0);
+	random_shuffle(all(v), [](ll x) { return ra() % x; });
+	Node* tr = new Node(v,0,N);
+	rep(i,0,N) rep(j,0,N) if (i <= j) {
+		ll ma = -inf;
+		rep(k,i,j) ma = max(ma, v[k]);
+		assert(ma == tr->query(i,j));
+	}
+	rep(it,0,1000000) {
+		ll i = ra() % (N+1), j = ra() % (N+1);
+		if (i > j) swap(i, j);
+		ll x = (ra() % 10) - 5;
+
+		ll r = ra() % 100;
+		if (r < 30) {
+			::res = tr->query(i, j);
+			ll ma = -inf;
+			rep(k,i,j) ma = max(ma, v[k]);
+			assert(ma == ::res);
+		}
+		else if (r < 70) {
+			tr->add(i, j, x);
+			rep(k,i,j) v[k] += x;
+		}
+		else {
+			tr->set(i, j, x);
+			rep(k,i,j) v[k] = x;
+		}
+	}
+	cout<<"Tests passed!"<<endl;
+}