generalise segment tree binary search

3 files changed, 66 insertions, 7 deletions

diff --git a/content/datastructures/datastructures.tex b/content/datastructures/datastructures.tex
index 4bbe507..0495543 100644
--- a/content/datastructures/datastructures.tex
+++ b/content/datastructures/datastructures.tex
@@ -10,7 +10,7 @@
 	
 	\subsubsection{Lazy Propagation}
 	Assignment modifications, sum queries \\
-	\method{lower\_bound}{erster Index in $[l, r)$ $\geq x$ (erfordert max-combine)}{\log(n)}
+	\method{binary\_search}{kleinstes $x$ in $[l, r]$ mit $f(\text{query}(l, x))$ (monoton in $x$)}{\log(n)}
 	\sourcecode{datastructures/lazyPropagation.cpp}
 \end{algorithm}
 
diff --git a/content/datastructures/lazyPropagation.cpp b/content/datastructures/lazyPropagation.cpp
index 82176cb..a5be822 100644
--- a/content/datastructures/lazyPropagation.cpp
+++ b/content/datastructures/lazyPropagation.cpp
@@ -63,21 +63,23 @@ struct SegTree {
 	}
 
 	// Optional:
-	int lower_bound(int l, int r, T x) {
+	int binary_search(int l, int r, auto &&f) {
+		if (f(E)) return l;
 		l += n, r += n;
-		push(l), push(r - 1);
+		push(l), push(r);
 		int a[64] = {}, lp = 0, rp = 64;
 		for (; l < r; l /= 2, r /= 2) {
 			if (l&1) a[lp++] = l++;
 			if (r&1) a[--rp] = --r;
 		}
-		for (int i : a) if (i != 0 && tree[i] >= x) { // Modify this
+		T x = E, y = x;
+		for (int i : a) if (i != 0 && f(x = comb(y = x, tree[i]))) {
 			while (i < n) {
 				push_down(i);
-				if (tree[2 * i] >= x) i = 2 * i; // And this
-				else i = 2 * i + 1;
+				if (f(x = comb(y, tree[2*i]))) i = 2 * i;
+				else i = 2 * i + 1, y = x;
 			}
-			return i - n;
+			return i - n + 1;
 		}
 		return -1;
 	}
diff --git a/test/datastructures/lazyPropagation.cpp b/test/datastructures/lazyPropagation.cpp
index 3030e1c..16db945 100644
--- a/test/datastructures/lazyPropagation.cpp
+++ b/test/datastructures/lazyPropagation.cpp
@@ -34,6 +34,39 @@ void stress_test() {
 	cerr << "tested random queries: " << queries << endl;
 }
 
+void stress_test_binary_search() {
+	ll queries = 0;
+	for (int tries = 0; tries < 100; tries++) {
+		int n = Random::integer<int>(10, 100);
+		vector<ll> naive = Random::integers<ll>(n, 0, 1000);
+		SegTree tree(naive);
+		for (int operations = 0; operations < 1000; operations++) {
+			{
+				int l = Random::integer<int>(0, n + 1);
+				int r = Random::integer<int>(0, n + 1);
+				//if (l > r) swap(l, r);
+				ll x = Random::integer<ll>(0, 1000);
+				tree.update(l, r, x);
+				for (int j = l; j < r; j++) naive[j] = x;
+			}
+			{
+				queries++;
+				int l = Random::integer<int>(0, n + 1);
+				int r = Random::integer<int>(0, n + 1);
+				ll x = Random::integer<ll>(0, 10000);
+				//if (l > r) swap(l, r);
+				int got = tree.binary_search(l, r, [x](ll v) { return v >= x; });
+				ll sum;
+				int j;
+				for (j = l, sum = 0; j < r && sum < x; j++) sum += naive[j];
+				int expected = sum >= x ? j : -1;
+				if (got != expected) cerr << "got: " << got << ", expected: " << expected << FAIL;
+			}
+		}
+	}
+	cerr << "tested random binary searches: " << queries << endl;
+}
+
 void performance_test() {
 	timer t;
 	t.start();
@@ -55,7 +88,31 @@ void performance_test() {
 	cerr << "tested performance: " << t.time << "ms (hash: " << hash << ")" << endl;
 }
 
+void performance_test_binary_search() {
+	timer t;
+	t.start();
+	vector<ll> tmp(N);
+	SegTree tree(tmp);
+	t.stop();
+	hash_t hash = 0;
+	for (int operations = 0; operations < N; operations++) {
+		auto [l1, r1] = Random::pair<int>(0, N + 1);
+		auto [l2, r2] = Random::pair<int>(0, N + 1);
+		ll x1 = Random::integer<ll>(0, 1000);
+		ll x2 = Random::integer<ll>(0, 1000 * N);
+
+		t.start();
+		tree.update(l1, r1, x1);
+		hash ^= tree.binary_search(l2, r2, [x2](ll v) { return v >= x2; });
+		t.stop();
+	}
+	if (t.time > 2000) cerr << "too slow: " << t.time << FAIL;
+	cerr << "tested performance: " << t.time << "ms (hash: " << hash << ")" << endl;
+}
+
 int main() {
 	stress_test();
+	stress_test_binary_search();
 	performance_test();
+	performance_test_binary_search();
 }