improve segment trees

3 files changed, 85 insertions, 87 deletions

diff --git a/datastructures/datastructures.tex b/datastructures/datastructures.tex
index 8f9698c..4139219 100644
--- a/datastructures/datastructures.tex
+++ b/datastructures/datastructures.tex
@@ -2,8 +2,8 @@
 
 \begin{algorithm}{Segmentbaum}
 	\begin{methods}
-		\method{init}{baut den Baum auf}{n}
-		\method{query}{findet das min(max) in [l, r)}{\log(n)}
+		\method{SegTree}{baut den Baum auf}{n}
+		\method{query}{findet Summe über [l, r)}{\log(n)}
 		\method{update}{ändert einen Wert}{\log(n)}
 	\end{methods}
 	\sourcecode{datastructures/segmentTree.cpp}
diff --git a/datastructures/lazyPropagation.cpp b/datastructures/lazyPropagation.cpp
index e8411fb..7af1a6a 100644
--- a/datastructures/lazyPropagation.cpp
+++ b/datastructures/lazyPropagation.cpp
@@ -1,86 +1,84 @@
-template<typename T = ll, typename U = ll>
 struct SegTree {
-	int size, height;
-	static constexpr T neutral = 0; // Neutral element for combine
-	static constexpr U updateFlag = 0; // Unused value by updates
+	using T = ll; using U = ll;
+	int n, h;
+	static constexpr T E = 0; // Neutral element for combine
+	static constexpr U UF = 0; // Unused value by updates
 	vector<T> tree; vector<U> lazy;
+	vector<ll> k; // size of segments (optional)
 
-	SegTree(const vector<T>& a) : SegTree(sz(a)) {
-		copy(all(a), tree.begin() + size);
-		for (int i = size - 1; i > 0; i--)
-			tree[i] = combine(tree[2 * i], tree[2 * i + 1]);
-	}
-
-	SegTree(int n) : size(n), height(__lg(2 * n)),
-		tree(2 * n, neutral), lazy(n, updateFlag) {}
+	SegTree(const vector<T>& a) : n(sz(a) + 1), tree(2 * n, E),
+	//SegTree(int size, T def = E) : n(size + 1), tree(2 * n, def),
+			h(__lg(2 * n)), lazy(n, UF), k(2 * n, 1) {
+		copy(all(a), tree.begin() + n);
+		for (int i = n - 1; i > 0; i--) {
+			k[i] = 2 * k[2 * i];
+			tree[i] = comb(tree[2 * i], tree[2 * i + 1]);
+	}}
 
-	T combine(T a, T b) {return a + b;} // Modify this + neutral
+	T comb(T a, T b) {return a + b;} // Modify this + E
 
-	void apply(int i, U val, int k) { // And this + updateFlag
-		tree[i] = val * k;
-		if (i < size) lazy[i] = val; // Don't forget this
+	void apply(int i, U val) { // And this + UF
+		tree[i] = val * k[i];
+		if (i < n) lazy[i] = val; // Don't forget this
 	}
 
-	void push_down(int i, int k) {
-		if (lazy[i] != updateFlag) {
-			apply(2 * i, lazy[i], k);
-			apply(2 * i + 1, lazy[i], k);
-			lazy[i] = updateFlag;
+	void push_down(int i) {
+		if (lazy[i] != UF) {
+			apply(2 * i, lazy[i]);
+			apply(2 * i + 1, lazy[i]);
+			lazy[i] = UF;
 	}}
 
 	void push(int i) {
-		for (int s = height, k = 1 << (height-1); s > 0; s--, k /= 2)
-			push_down(i >> s, k);
+		for (int s = h; s > 0; s--) push_down(i >> s);
 	}
 
 	void build(int i) {
-		for (int k = 2; i /= 2; k *= 2) {
-			push_down(i, k / 2);
-			tree[i] = combine(tree[2 * i], tree[2 * i + 1]);
+		while (i /= 2) {
+			push_down(i);
+			tree[i] = comb(tree[2 * i], tree[2 * i + 1]);
 	}}
 
-	void update(int l, int r, U val) { // data[l..r) = val
-		l += size, r += size;
+	void update(int l, int r, U val) {
+		l += n, r += n;
 		int l0 = l, r0 = r;
 		push(l0), push(r0 - 1);
-		for (int k = 1; l < r; l /= 2, r /= 2, k *= 2) {
-			if (l&1) apply(l++, val, k);
-			if (r&1) apply(--r, val, k);
+		for (; l < r; l /= 2, r /= 2) {
+			if (l&1) apply(l++, val);
+			if (r&1) apply(--r, val);
 		}
 		build(l0), build(r0 - 1);
 	}
 
-	T query(int l, int r) { // sum[l..r)
-		l += size, r += size;
+	T query(int l, int r) {
+		l += n, r += n;
 		push(l), push(r - 1);
-		T resL = neutral, resR = neutral;
+		T resL = E, resR = E;
 		for (; l < r; l /= 2, r /= 2) {
-			if (l&1) resL = combine(resL, tree[l++]);
-			if (r&1) resR = combine(tree[--r], resR);
+			if (l&1) resL = comb(resL, tree[l++]);
+			if (r&1) resR = comb(tree[--r], resR);
 		}
-		return combine(resL, resR);
+		return comb(resL, resR);
 	}
 
 	// Optional:
 	ll lower_bound(int l, int r, T x) {
-		l += size, r += size;
+		l += n, r += n;
 		push(l), push(r - 1);
-		vector<pair<int, int>> a, st;
-		for (int k = 1; l < r; l /= 2, r /= 2, k *= 2) {
-			if (l&1) a.emplace_back(l++, k);
-			if (r&1) st.emplace_back(--r, k);
+		vector<int> a, st;
+		for (; l < r; l /= 2, r /= 2) {
+			if (l&1) a.push_back(l++);
+			if (r&1) st.push_back(--r);
 		}
 		a.insert(a.end(), st.rbegin(), st.rend());
-		for (auto [i, k] : a) {
-			if (tree[i] >= x) return find(i, x, k); // Modify this
+		for (int i : a) if (tree[i] >= x) { // Modify this
+			while (i < n) {
+				push_down(i);
+				if (tree[2 * i] >= x) i = 2 * i; // And this
+				else i = 2 * i + 1;
+			}
+			return i - n;
 		}
 		return -1;
 	}
-
-	ll find(int i, T x, int k) {
-		if (i >= size) return i - size;
-		push_down(i, k / 2);
-		if (tree[2*i] >= x) return find(2 * i, x, k / 2); // And this
-		else return find(2 * i + 1, x, k / 2);
-	}
 };
diff --git a/datastructures/segmentTree.cpp b/datastructures/segmentTree.cpp
index 7b97f57..5a75d69 100644
--- a/datastructures/segmentTree.cpp
+++ b/datastructures/segmentTree.cpp
@@ -1,42 +1,42 @@
-vector<ll> tree;
-constexpr ll neutral = 0; // Neutral element for combine
+struct SegTree {
+	using T = ll;
+	int n;
+	vector<T> tree;
+	static constexpr T E = 0; // Neutral element for combine
 
-ll combine(ll a, ll b) {
-	return a + b;
-}
+	SegTree(vector<T>& a) : n(sz(a)), tree(2 * n) {
+	//SegTree(int size, T val = E) : n(size), tree(2 * n, val) {
+		copy(all(a), tree.begin() + n);
+		for (int i = n - 1; i > 0; i--) { // remove for range update
+			tree[i] = comb(tree[2 * i], tree[2 * i + 1]);
+	}}
 
-void init(vector<ll>& a) {
-	tree.assign(2 * sz(a), 0);
-	copy(all(a), tree.begin() + sz(a));
-	for (int i = sz(tree)/2 - 1; i > 0; i--) {
-		tree[i] = combine(tree[2 * i], tree[2 * i + 1]);
-}}
+	ll comb(T a, T b) { return a + b; } // modify this + neutral
 
-void update(int i, ll val) {
-	for (tree[i += sz(tree)/2] = val; i /= 2; ) {
-		tree[i] = combine(tree[2 * i], tree[2 * i + 1]);
-}}
+	void update(int i, T val) {
+		tree[i += n] = val; // apply update code
+		while (i /= 2) tree[i] = comb(tree[2 * i], tree[2 * i + 1]);
+	}
 
-ll query(int l, int r) {
-	ll resL = neutral, resR = neutral;
-	for (l += sz(tree)/2, r += sz(tree)/2; l < r; l /= 2, r /= 2) {
-		if (l&1) resL = combine(resL, tree[l++]);
-		if (r&1) resR = combine(tree[--r], resR);
+	T query(int l, int r) {
+		T resL = E, resR = E;
+		for (l += n, r += n; l < r; l /= 2, r /= 2) {
+			if (l&1) resL = comb(resL, tree[l++]);
+			if (r&1) resR = comb(tree[--r], resR);
+		}
+		return comb(resL, resR);
 	}
-	return combine(resL, resR);
-}
 
-// Oder: Intervall-Modifikation, Punkt-Query:
-void modify(int l, int r, ll val) {
-	for (l += sz(tree)/2, r += sz(tree)/2; l < r; l /= 2, r /= 2) {
-		if (l&1) {tree[l] = combine(tree[l], val); l++;};
-		if (r&1) {--r; tree[r] = combine(tree[r], val);};
-}}
+	// OR: range update + point query, needs commutative comb
+	void modify(int l, int r, T val) {
+		for (l += n, r += n; l < r; l /= 2, r /= 2) {
+			if (l&1) tree[l] = comb(tree[l], val), l++;
+			if (r&1) --r, tree[r] = comb(tree[r], val);
+	}}
 
-ll query(int i) {
-	ll res = neutral;
-	for (i += sz(tree)/2; i > 0; i /= 2) {
-		res = combine(res, tree[i]);
+	T query(int i) {
+		T res = E;
+		for (i += n; i > 0; i /= 2) res = comb(res, tree[i]);
+		return res;
 	}
-	return res;
-}
+};