Skip to content

Commit d9ef28b

Browse files
Mr-Python-in-ChinaMr-Python-in-China
committed
enable merge function of segtree to get zhe node size; release 0.5.0
1 parent 8ef189c commit d9ef28b

File tree

10 files changed

+686
-247
lines changed

10 files changed

+686
-247
lines changed

‎library/example/Luogu_P3372.cpp‎

Lines changed: 133 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
// mp-oi-library extension note: start. Do not modify this part.
2-
// mp-oi-library extension header: library version: @MrPython-0.2.6
2+
// mp-oi-library extension header: library version: @MrPython-0.4.2
33
// mp-oi-library extension note: file bits.hpp
44
#ifndef MP_LIBRARY_BITS
55
#define MP_LIBRARY_BITS
@@ -32,6 +32,30 @@ unsigned long long highbit(unsigned long long x) {
3232
}
3333
} // namespace mrpython
3434
#endif
35+
// mp-oi-library extension note: file segment_tree_utility.hpp
36+
#ifndef MP_LIBRARY_SEGMENT_TREE_UTILITY_HPP
37+
#define MP_LIBRARY_SEGMENT_TREE_UTILITY_HPP
38+
#include <cstddef>
39+
#include <utility>
40+
namespace mrpython {
41+
template <typename F, typename T>
42+
auto call_merge_impl(F const& f, T const& lv, std::size_t ls, T const& rv,
43+
std::size_t rs, int) -> decltype(f(lv, ls, rv, rs)) {
44+
return f(lv, ls, rv, rs);
45+
}
46+
template <typename F, typename T>
47+
auto call_merge_impl(F const& f, T const& lv, std::size_t, T const& rv,
48+
std::size_t, long) -> decltype(f(lv, rv)) {
49+
return f(lv, rv);
50+
}
51+
template <typename F, typename T>
52+
decltype(auto) call_merge(F const& f, T const& lv, std::size_t ls, T const& rv,
53+
std::size_t rs) {
54+
return call_merge_impl(f, lv, ls, rv, rs, 0);
55+
}
56+
} // namespace mrpython
57+
#endif // MP_LIBRARY_SEGMENT_TREE_UTILITY_HPP
58+
3559
// mp-oi-library extension note: file lazy_segment_tree.hpp
3660

3761
#ifndef MP_LIBRARY_LAZY_SEGMENT_TREE_HPP
@@ -53,12 +77,22 @@ class lazy_segment_tree {
5377
MergeFunction mergeData;
5478
OperateFunction operate;
5579
MergeLazyFunction mergeLazy;
80+
auto doMergeData(std::pair<T const&, size_t> l,
81+
std::pair<T const&, size_t> r) {
82+
return std::make_pair(
83+
call_merge(mergeData, l.first, l.second, r.first, r.second),
84+
l.second + r.second);
85+
}
86+
decltype(auto) doMergeData(T const& lv, size_t ls, T const& rv, size_t rs) {
87+
return call_merge(mergeData, lv, ls, rv, rs);
88+
}
5689
void build(void) {
5790
data.reserve(2 * n - 1), size.reserve(2 * n - 1);
5891
for (size_t i = n; i < 2 * n - 1; ++i) {
5992
size_t d = 2 * n - 1 - i;
6093
size_t l = d * 2, r = d * 2 + 1;
61-
data.emplace_back(mergeData(data[2 * n - 1 - l], data[2 * n - 1 - r]));
94+
data.emplace_back(doMergeData(data[2 * n - 1 - l], size[2 * n - 1 - l],
95+
data[2 * n - 1 - r], size[2 * n - 1 - r]));
6296
size.emplace_back(size[2 * n - 1 - l] + size[2 * n - 1 - r]);
6397
}
6498
std::reverse(data.begin(), data.end());
@@ -77,13 +111,13 @@ class lazy_segment_tree {
77111
data[pos] = operate(lazyVal, data[pos], size[pos]);
78112
lazy[pos] = mergeLazy(lazy[pos], lazyVal);
79113
}
80-
T get_impl(size_t l, size_t r, size_t pos) {
81-
if (l == 0 && r == size[pos]) return data[pos];
114+
std::pair<T, size_t> get_impl(size_t l, size_t r, size_t pos) {
115+
if (l == 0 && r == size[pos]) return {data[pos], size[pos]};
82116
pushdown(pos);
83117
size_t m = size[pos * 2 + 1];
84118
if (l < m && r > m)
85-
return mergeData(get_impl(l, m, pos * 2 + 1),
86-
get_impl(0, r - m, pos * 2 + 2));
119+
return doMergeData(get_impl(l, m, pos * 2 + 1),
120+
get_impl(0, r - m, pos * 2 + 2));
87121
else if (l < m)
88122
return get_impl(l, r, pos * 2 + 1);
89123
else if (r > m)
@@ -104,7 +138,35 @@ class lazy_segment_tree {
104138
set_impl(l - m, r - m, operateVal, pos * 2 + 2);
105139
else
106140
__builtin_unreachable();
107-
data[pos] = mergeData(data[pos * 2 + 1], data[pos * 2 + 2]);
141+
data[pos] = doMergeData(data[pos * 2 + 1], size[pos * 2 + 1],
142+
data[pos * 2 + 2], size[pos * 2 + 2]);
143+
}
144+
template <typename Operate>
145+
void set_impl(size_t c, Operate const& opv, size_t pos) {
146+
if (size[pos] == 1) {
147+
data[pos] = opv((T const&)data[pos]);
148+
return;
149+
}
150+
pushdown(pos);
151+
size_t m = size[pos * 2 + 1];
152+
if (c < m)
153+
set_impl(c, opv, pos * 2 + 1);
154+
else
155+
set_impl(c - m, opv, pos * 2 + 2);
156+
data[pos] = doMergeData(data[pos * 2 + 1], size[pos * 2 + 1],
157+
data[pos * 2 + 2], size[pos * 2 + 2]);
158+
}
159+
size_t data_id_to_node_id(size_t x) {
160+
x += n - ((2 * n - 1) - (highbit(2 * n - 1) - 1));
161+
if (x >= n) x -= n;
162+
x += n - 1;
163+
return x;
164+
}
165+
size_t node_id_to_data_id(size_t x) {
166+
x -= n - 1;
167+
x += ((2 * n - 1) - (highbit(2 * n - 1) - 1));
168+
if (x >= n) x -= n;
169+
return x;
108170
}
109171

110172
public:
@@ -122,8 +184,8 @@ class lazy_segment_tree {
122184
mergeData(mergeDataFun),
123185
operate(OperateFun),
124186
mergeLazy(mergeTagFun) {
125-
rotate(data.begin(), data.begin() + (2 * n - 1) - (highbit(2 * n - 1) - 1),
126-
data.end());
187+
rotate(data.begin(),
188+
data.begin() + ((2 * n - 1) - (highbit(2 * n - 1) - 1)), data.end());
127189
reverse(data.begin(), data.end());
128190
build();
129191
}
@@ -141,10 +203,71 @@ class lazy_segment_tree {
141203
mergeLazy(mergeTagFun) {
142204
build();
143205
}
144-
T get(size_t l, size_t r) { return get_impl(l, r, 0); }
206+
T get(size_t l, size_t r) { return get_impl(l, r, 0).first; }
207+
T getd(size_t l, size_t r, T const& e = {}) { return l == r ? e : get(l, r); }
145208
void set(size_t l, size_t r, Lazy const& operateVal) {
209+
if (l == r) return;
146210
set_impl(l, r, operateVal, 0);
147211
}
212+
template <typename Operate> void set(size_t p, Operate const& opv) {
213+
set_impl(p, opv, 0);
214+
}
215+
template <typename Check>
216+
size_t find_first_right(size_t l, Check const& check) {
217+
if (l >= n) return l;
218+
set(l, [](T const& _) { return _; }); // 将一路上的节点都 pushdown 下去
219+
l = data_id_to_node_id(l);
220+
while (l % 2 == 1) l /= 2;
221+
while (l < 2 * n - 1 && check(data[l])) l = l * 2 + 1;
222+
if (l >= 2 * n - 1) return node_id_to_data_id(l / 2);
223+
std::pair<T, size_t> v = {data[l], size[l]};
224+
while (true) {
225+
++l;
226+
if (!(l & (l + 1))) return n;
227+
while (l % 2 == 1) l /= 2;
228+
std::pair<T, size_t> vl = doMergeData(v, {data[l], size[l]});
229+
if (check(vl.first)) break;
230+
v = std::move(vl);
231+
}
232+
while (l < n - 1) {
233+
pushdown(l);
234+
std::pair<T, size_t> vl =
235+
doMergeData(v, {data[l * 2 + 1], size[l * 2 + 1]});
236+
if (!check(vl.first))
237+
l = l * 2 + 2, v = std::move(vl);
238+
else
239+
l = l * 2 + 1;
240+
}
241+
return node_id_to_data_id(l);
242+
}
243+
template <typename Check>
244+
size_t find_last_left(size_t r, Check const& check) {
245+
if (r >= n) return r;
246+
set(r, [](T const& _) { return _; }); // 将一路上的节点都 pushdown 下去
247+
r = data_id_to_node_id(r);
248+
while (r && r % 2 == 0) r = (r - 1) / 2;
249+
while (r < 2 * n - 1 && check(data[r])) r = r * 2 + 2;
250+
if (r >= 2 * n - 1) return node_id_to_data_id((r - 1) / 2);
251+
std::pair<T, size_t> v = {data[r], size[r]};
252+
while (true) {
253+
if (!(r & (r + 1))) return -1;
254+
--r;
255+
while (r % 2 == 0) r = (r - 1) / 2;
256+
std::pair<T, size_t> vl = doMergeData(v, {data[r], size[r]});
257+
if (check(vl.first)) break;
258+
v = std::move(vl);
259+
}
260+
while (r < n - 1) {
261+
pushdown(r);
262+
std::pair<T, size_t> vr =
263+
doMergeData(v, {data[r * 2 + 2], size[r * 2 + 2]});
264+
if (!check(vr.first))
265+
r = r * 2 + 1, v = std::move(vr);
266+
else
267+
r = r * 2 + 2;
268+
}
269+
return node_id_to_data_id(r);
270+
}
148271
};
149272
template <typename T> struct lazy_segment_tree_add_add_operate_function {
150273
T operator()(T const& lazy, T const& data, size_t size) const {

0 commit comments

Comments
 (0)