diff --git a/library/data_structures_[l,r)/seg_tree.hpp b/library/data_structures_[l,r)/seg_tree.hpp index be34b3bf..67de2901 100644 --- a/library/data_structures_[l,r)/seg_tree.hpp +++ b/library/data_structures_[l,r)/seg_tree.hpp @@ -5,6 +5,9 @@ //! tree st2(n, INT_MAX, [&](int x, int y) -> int { //! return min(x, y); //! }); +//! int idx = st2.walk(l, r, [&](int value) { +//! return value > x; +//! }); // min idx in [l, r) such that f is false //! @endcode //! @time O(n + q log n) //! @space O(n) @@ -27,4 +30,5 @@ template struct tree { } return op(x, y); } +#include "seg_tree_uncommon/walk.hpp" }; diff --git a/library/data_structures_[l,r)/seg_tree_uncommon/walk.hpp b/library/data_structures_[l,r)/seg_tree_uncommon/walk.hpp new file mode 100644 index 00000000..ecc72f69 --- /dev/null +++ b/library/data_structures_[l,r)/seg_tree_uncommon/walk.hpp @@ -0,0 +1,8 @@ +int walk(int l, int r, const auto& f) { + while (l < r) { + int u = l + n, x = __lg(min(u & -u, r - l)); + if (f(s[u >> x])) l += 1 << x; + else r = l + (1 << x) - 1; + } + return l; +} diff --git a/tests/library_checker_aizu_tests/data_structures/simple_tree_walk.test.cpp b/tests/library_checker_aizu_tests/data_structures/simple_tree_walk.test.cpp new file mode 100644 index 00000000..8155a78e --- /dev/null +++ b/tests/library_checker_aizu_tests/data_structures/simple_tree_walk.test.cpp @@ -0,0 +1,44 @@ +#define PROBLEM \ + "https://judge.yosupo.jp/problem/predecessor_problem" +#include "../template.hpp" +#include "../../../library/data_structures_[l,r)/seg_tree.hpp" +int main() { + cin.tie(0)->sync_with_stdio(0); + int n, q; + cin >> n >> q; + tree st(n, 0, plus{}); + string s; + cin >> s; + for (int i = 0; i < n; i++) st.update(i, s[i] - '0'); + while (q--) { + int type, k; + cin >> type >> k; + if (type == 0) { + if (st.query(k, k + 1) == 0) st.update(k, 1); + } else if (type == 1) { + if (st.query(k, k + 1) == 1) st.update(k, 0); + } else if (type == 2) { + cout << st.query(k, k + 1) << '\n'; + } else if (type == 3) { + // returns first element in [k,n) such that sum > 0 + int idx = + st.walk(k, n, [&](int sum) { return sum == 0; }); + if (idx == n) idx = -1; + cout << idx << '\n'; + } else { + assert(type == 4); + int total = st.query(0, k + 1); + if (total == 0) { + cout << -1 << '\n'; + } else { + int pref_sum = 0; + cout << st.walk(0, k + 1, [&](int sum) { + if (pref_sum + sum < total) + return pref_sum += sum, 1; + return 0; + }) << '\n'; + } + } + } + return 0; +}