Skip to content

Commit d4c8718

Browse files
Mr-Python-in-ChinaMr-Python-in-China
committed
Support find_last_left and find_first_right for lazy_segment_tree; fix bug for typical_segment_tree.
1 parent 92d9997 commit d4c8718

File tree

8 files changed

+242
-53
lines changed

8 files changed

+242
-53
lines changed

‎library/mrpython/lazy_segment_tree.hpp

Lines changed: 87 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,32 @@ class lazy_segment_tree {
7272
__builtin_unreachable();
7373
data[pos] = mergeData(data[pos * 2 + 1], data[pos * 2 + 2]);
7474
}
75+
template <typename Operate>
76+
void set_impl(size_t c, Operate const& opv, size_t pos) {
77+
if (size[pos] == 1) {
78+
data[pos] = opv((T const&)data[pos]);
79+
return;
80+
}
81+
pushdown(pos);
82+
size_t m = size[pos * 2 + 1];
83+
if (c < m)
84+
set_impl(c, opv, pos * 2 + 1);
85+
else
86+
set_impl(c - m, opv, pos * 2 + 2);
87+
data[pos] = mergeData(data[pos * 2 + 1], data[pos * 2 + 2]);
88+
}
89+
size_t data_id_to_node_id(size_t x) {
90+
x += n - ((2 * n - 1) - (highbit(2 * n - 1) - 1));
91+
if (x >= n) x -= n;
92+
x += n - 1;
93+
return x;
94+
}
95+
size_t node_id_to_data_id(size_t x) {
96+
x -= n - 1;
97+
x += ((2 * n - 1) - (highbit(2 * n - 1) - 1));
98+
if (x >= n) x -= n;
99+
return x;
100+
}
75101

76102
public:
77103
template <typename InputIterator>
@@ -88,8 +114,8 @@ class lazy_segment_tree {
88114
mergeData(mergeDataFun),
89115
operate(OperateFun),
90116
mergeLazy(mergeTagFun) {
91-
rotate(data.begin(), data.begin() + (2 * n - 1) - (highbit(2 * n - 1) - 1),
92-
data.end());
117+
rotate(data.begin(),
118+
data.begin() + ((2 * n - 1) - (highbit(2 * n - 1) - 1)), data.end());
93119
reverse(data.begin(), data.end());
94120
build();
95121
}
@@ -108,9 +134,68 @@ class lazy_segment_tree {
108134
build();
109135
}
110136
T get(size_t l, size_t r) { return get_impl(l, r, 0); }
137+
T getd(size_t l, size_t r, T const& e = {}) { return l == r ? e : get(l, r); }
111138
void set(size_t l, size_t r, Lazy const& operateVal) {
139+
if (l == r) return;
112140
set_impl(l, r, operateVal, 0);
113141
}
142+
template <typename Operate> void set(size_t p, Operate const& opv) {
143+
set_impl(p, opv, 0);
144+
}
145+
template <typename Check>
146+
size_t find_first_right(size_t l, Check const& check) {
147+
if (l >= n) return l;
148+
set(l, [](T const& _) { return _; });
149+
l = data_id_to_node_id(l);
150+
while (l % 2 == 1) l /= 2;
151+
while (l < 2 * n - 1 && check(data[l])) l = l * 2 + 1;
152+
if (l >= 2 * n - 1) return node_id_to_data_id(l / 2);
153+
T v = data[l];
154+
while (true) {
155+
++l;
156+
if (!(l & (l + 1))) return n;
157+
while (l % 2 == 1) l /= 2;
158+
T vl = mergeData(v, data[l]);
159+
if (check(vl)) break;
160+
v = vl;
161+
}
162+
while (l < n - 1) {
163+
pushdown(l);
164+
T vl = mergeData(v, data[l * 2 + 1]);
165+
if (!check(vl))
166+
l = l * 2 + 2, v = vl;
167+
else
168+
l = l * 2 + 1;
169+
}
170+
return node_id_to_data_id(l);
171+
}
172+
template <typename Check>
173+
size_t find_last_left(size_t r, Check const& check) {
174+
if (r >= n) return r;
175+
set(r, [](T const& _) { return _; });
176+
r = data_id_to_node_id(r);
177+
while (r && r % 2 == 0) r = (r - 1) / 2;
178+
while (r < 2 * n - 1 && check(data[r])) r = r * 2 + 2;
179+
if (r >= 2 * n - 1) return node_id_to_data_id((r - 1) / 2);
180+
T v = data[r];
181+
while (true) {
182+
if (!(r & (r + 1))) return -1;
183+
--r;
184+
while (r % 2 == 0) r = (r - 1) / 2;
185+
T vl = mergeData(v, data[r]);
186+
if (check(vl)) break;
187+
v = vl;
188+
}
189+
while (r < n - 1) {
190+
pushdown(r);
191+
T vr = mergeData(v, data[r * 2 + 2]);
192+
if (!check(vr))
193+
r = r * 2 + 1, v = vr;
194+
else
195+
r = r * 2 + 2;
196+
}
197+
return node_id_to_data_id(r);
198+
}
114199
};
115200
template <typename T> struct lazy_segment_tree_add_add_operate_function {
116201
T operator()(T const& lazy, T const& data, size_t size) const {

‎library/mrpython/typical_segment_tree.hpp

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -107,11 +107,14 @@ template <typename T, typename MergeFunction> class typical_segment_tree {
107107
while (l < 2 * n - 1 && check(data[l])) l = l * 2 + 1;
108108
if (l >= 2 * n - 1) return node_id_to_data_id(l / 2);
109109
T v = data[l];
110-
do {
110+
while (true) {
111111
++l;
112112
if (!(l & (l + 1))) return n;
113113
while (l % 2 == 1) l /= 2;
114-
} while (!check(data[l]));
114+
T vl = merge(v, data[l]);
115+
if (check(vl)) break;
116+
v = vl;
117+
}
115118
while (l < n - 1) {
116119
T vl = merge(v, data[l * 2 + 1]);
117120
if (!check(vl))
@@ -129,11 +132,14 @@ template <typename T, typename MergeFunction> class typical_segment_tree {
129132
while (r < 2 * n - 1 && check(data[r])) r = r * 2 + 2;
130133
if (r >= 2 * n - 1) return node_id_to_data_id((r - 1) / 2);
131134
T v = data[r];
132-
do {
135+
while (true) {
133136
if (!(r & (r + 1))) return -1;
134137
--r;
135138
while (r % 2 == 0) r = (r - 1) / 2;
136-
} while (!check(data[r]));
139+
T vl = merge(v, data[r]);
140+
if (check(vl)) break;
141+
v = vl;
142+
}
137143
while (r < n - 1) {
138144
T vr = merge(v, data[r * 2 + 2]);
139145
if (!check(vr))

‎library/test/lazy_segment_tree/all.hpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
11
#include "lazy_segment_tree_add_add.cpp"
2+
#include "lazy_segment_tree_binary_search.cpp"

‎library/test/lazy_segment_tree/lazy_segment_tree_add_add.cpp

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#include <gtest/gtest.h>
22

3+
#include <cstddef>
34
#include <mrpython/lazy_segment_tree.hpp>
45
#include <random>
56

@@ -10,22 +11,31 @@ TEST(lazy_segment_tree, add) {
1011
std::uniform_int_distribution<unsigned> val_dist(
1112
std::numeric_limits<unsigned>::min(),
1213
std::numeric_limits<unsigned>::max()),
13-
size_dist(0, n - 1), operator_dist(0, 1);
14+
size_dist(0, n - 1), operator_dist(0, 2);
1415
std::vector<unsigned> a(n);
1516
std::generate(a.begin(), a.end(), [&] { return val_dist(gen); });
1617
mrpython::lazy_segment_tree_add_add<unsigned> tree(a.begin(), a.end(), 0);
1718
while (q--) {
18-
std::size_t l = size_dist(gen), r = size_dist(gen);
19-
if (l > r) std::swap(l, r);
20-
assert(l < r + 1);
21-
if (operator_dist(gen)) {
19+
unsigned char op = operator_dist(gen);
20+
if (op == 0) {
21+
std::size_t l = size_dist(gen), r = size_dist(gen);
22+
if (l > r) std::swap(l, r);
23+
assert(l < r + 1);
2224
unsigned ans =
2325
std::accumulate(a.begin() + l, a.begin() + r + 1, (unsigned)0);
2426
EXPECT_EQ(tree.get(l, r + 1), ans);
25-
} else {
27+
} else if (op == 1) {
28+
std::size_t l = size_dist(gen), r = size_dist(gen);
29+
if (l > r) std::swap(l, r);
30+
assert(l < r + 1);
2631
unsigned value = val_dist(gen);
2732
for (std::size_t i = l; i < r + 1; ++i) a[i] += value;
2833
tree.set(l, r + 1, value);
34+
} else {
35+
size_t p = size_dist(gen);
36+
unsigned value = val_dist(gen);
37+
a[p] += value;
38+
tree.set(p, [&](unsigned x) { return x + value; });
2939
}
3040
}
3141
}

‎library/test/lazy_segment_tree/lazy_segment_tree_binary_search.cpp

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
#include <gtest/gtest.h>
2+
3+
#include <cstddef>
4+
#include <mrpython/lazy_segment_tree.hpp>
5+
#include <random>
6+
7+
TEST(lazy_segment_tree, find_first_right) {
8+
std::mt19937_64 gen(std::random_device{}());
9+
std::size_t n = std::uniform_int_distribution<std::size_t>{1, 5000}(gen),
10+
q = std::uniform_int_distribution<std::size_t>{1, 5000}(gen);
11+
std::uniform_int_distribution<unsigned> val_dist(
12+
std::numeric_limits<unsigned>::min(),
13+
std::numeric_limits<unsigned>::max()),
14+
size_dist(0, n - 1), operator_dist(0, 1);
15+
std::uniform_int_distribution<unsigned long long> q_dist(
16+
std::numeric_limits<unsigned>::min(),
17+
(unsigned long long)std::numeric_limits<unsigned>::max() * n);
18+
std::vector<unsigned> a(n);
19+
std::generate(a.begin(), a.end(), [&] { return val_dist(gen); });
20+
mrpython::lazy_segment_tree_add_add<unsigned long long> tree(a.begin(),
21+
a.end(), 0);
22+
while (q--) {
23+
if (operator_dist(gen)) {
24+
std::size_t l = size_dist(gen), r = size_dist(gen);
25+
if (l > r) std::swap(l, r);
26+
assert(l < r + 1);
27+
unsigned long long ans =
28+
std::accumulate(a.begin() + l, a.begin() + r + 1, 0ull);
29+
EXPECT_EQ(tree.get(l, r + 1), ans);
30+
} else {
31+
std::size_t l = size_dist(gen), v = q_dist(gen);
32+
size_t p = l;
33+
unsigned long long s = 0;
34+
while (p < n) {
35+
s += a[p];
36+
if (s > v) break;
37+
++p;
38+
}
39+
size_t out=tree.find_first_right(l, [&](unsigned long long x) { return x > v; });
40+
EXPECT_EQ(out, p);
41+
}
42+
}
43+
}
44+
45+
TEST(lazy_segment_tree, find_last_left) {
46+
std::mt19937_64 gen(std::random_device{}());
47+
std::size_t n = std::uniform_int_distribution<std::size_t>{1, 5000}(gen),
48+
q = std::uniform_int_distribution<std::size_t>{1, 5000}(gen);
49+
std::uniform_int_distribution<unsigned> val_dist(
50+
std::numeric_limits<unsigned>::min(),
51+
std::numeric_limits<unsigned>::max()),
52+
size_dist(0, n - 1), operator_dist(0, 1);
53+
std::uniform_int_distribution<unsigned long long> q_dist(
54+
std::numeric_limits<unsigned>::min(),
55+
(unsigned long long)std::numeric_limits<unsigned>::max() * n);
56+
std::vector<unsigned> a(n);
57+
std::generate(a.begin(), a.end(), [&] { return val_dist(gen); });
58+
mrpython::lazy_segment_tree_add_add<unsigned long long> tree(a.begin(),
59+
a.end(), 0);
60+
while (q--) {
61+
if (operator_dist(gen)) {
62+
std::size_t l = size_dist(gen), r = size_dist(gen);
63+
if (l > r) std::swap(l, r);
64+
assert(l < r + 1);
65+
unsigned long long ans =
66+
std::accumulate(a.begin() + l, a.begin() + r + 1, 0ull);
67+
EXPECT_EQ(tree.get(l, r + 1), ans);
68+
} else {
69+
std::size_t l = size_dist(gen), v = q_dist(gen);
70+
size_t p = l;
71+
unsigned long long s = 0;
72+
while (p < n) {
73+
s += a[p];
74+
if (s > v) break;
75+
--p;
76+
}
77+
size_t out=tree.find_last_left(l, [&](unsigned long long x) { return x > v; });
78+
EXPECT_EQ(out, p);
79+
}
80+
}
81+
}

‎library/test/typical_segment_tree/all.hpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
11
#include "typical_segment_tree_add.cpp"
22
#include "typical_segment_tree_max.cpp"
33
#include "typical_segment_tree_min.cpp"
4+
#include "typical_segment_tree_binary_search.cpp"

‎library/test/typical_segment_tree/typical_segment_tree_binary_search.cpp

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
#include <gtest/gtest.h>
2+
3+
#include <cstddef>
4+
#include <mrpython/typical_segment_tree.hpp>
5+
#include <random>
6+
7+
TEST(typicaL_segment_tree, find_first_right) {
8+
std::mt19937_64 gen(std::random_device{}());
9+
std::size_t n = std::uniform_int_distribution<std::size_t>{1, 5000}(gen),
10+
q = std::uniform_int_distribution<std::size_t>{1, 5000}(gen);
11+
std::vector<int> a(n);
12+
std::uniform_int_distribution<int> val_dist(std::numeric_limits<int>::min(),
13+
std::numeric_limits<int>::max());
14+
std::uniform_int_distribution<std::size_t> size_dist(0, n),
15+
operator_dist(0, 1);
16+
std::generate(a.begin(), a.end(), [&] { return val_dist(gen); });
17+
mrpython::typical_segment_tree_min<int> tree(a.begin(), a.end());
18+
while (q--) {
19+
std::size_t l = size_dist(gen);
20+
int tv = val_dist(gen);
21+
auto check = [tv](int x) { return x <= tv; };
22+
size_t ans = std::find_if(a.begin() + l, a.end(), check) - a.begin();
23+
EXPECT_EQ(tree.find_first_right(l, check), ans);
24+
}
25+
}
26+
27+
TEST(typicaL_segment_tree, find_last_left) {
28+
std::mt19937_64 gen(std::random_device{}());
29+
std::size_t n = std::uniform_int_distribution<std::size_t>{1, 5000}(gen),
30+
q = std::uniform_int_distribution<std::size_t>{1, 5000}(gen);
31+
std::vector<int> a(n);
32+
std::uniform_int_distribution<int> val_dist(std::numeric_limits<int>::min(),
33+
std::numeric_limits<int>::max());
34+
std::uniform_int_distribution<std::size_t> size_dist(0, n),
35+
operator_dist(0, 1);
36+
std::generate(a.begin(), a.end(), [&] { return val_dist(gen); });
37+
mrpython::typical_segment_tree_min<int> tree(a.begin(), a.end());
38+
while (q--) {
39+
std::size_t l = size_dist(gen) - 1;
40+
int tv = val_dist(gen);
41+
auto check = [tv](int x) { return x <= tv; };
42+
size_t ans = a.rend() - std::find_if(a.rend() - l - 1, a.rend(), check) - 1;
43+
size_t out = tree.find_last_left(l, check);
44+
EXPECT_EQ(out, ans);
45+
}
46+
}

‎library/test/typical_segment_tree/typical_segment_tree_min.cpp

Lines changed: 0 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -29,44 +29,3 @@ TEST(typical_segment_tree, min) {
2929
}
3030
}
3131
}
32-
33-
TEST(typicaL_segment_tree, find_first_right) {
34-
std::mt19937_64 gen(std::random_device{}());
35-
std::size_t n = std::uniform_int_distribution<std::size_t>{1, 5000}(gen),
36-
q = std::uniform_int_distribution<std::size_t>{1, 5000}(gen);
37-
std::vector<int> a(n);
38-
std::uniform_int_distribution<int> val_dist(std::numeric_limits<int>::min(),
39-
std::numeric_limits<int>::max());
40-
std::uniform_int_distribution<std::size_t> size_dist(0, n),
41-
operator_dist(0, 1);
42-
std::generate(a.begin(), a.end(), [&] { return val_dist(gen); });
43-
mrpython::typical_segment_tree_min<int> tree(a.begin(), a.end());
44-
while (q--) {
45-
std::size_t l = size_dist(gen);
46-
int tv = val_dist(gen);
47-
auto check = [tv](int x) { return x <= tv; };
48-
size_t ans = std::find_if(a.begin() + l, a.end(), check) - a.begin();
49-
EXPECT_EQ(tree.find_first_right(l, check), ans);
50-
}
51-
}
52-
53-
TEST(typicaL_segment_tree, find_last_left) {
54-
std::mt19937_64 gen(std::random_device{}());
55-
std::size_t n = std::uniform_int_distribution<std::size_t>{1, 5000}(gen),
56-
q = std::uniform_int_distribution<std::size_t>{1, 5000}(gen);
57-
std::vector<int> a(n);
58-
std::uniform_int_distribution<int> val_dist(std::numeric_limits<int>::min(),
59-
std::numeric_limits<int>::max());
60-
std::uniform_int_distribution<std::size_t> size_dist(0, n),
61-
operator_dist(0, 1);
62-
std::generate(a.begin(), a.end(), [&] { return val_dist(gen); });
63-
mrpython::typical_segment_tree_min<int> tree(a.begin(), a.end());
64-
while (q--) {
65-
std::size_t l = size_dist(gen) - 1;
66-
int tv = val_dist(gen);
67-
auto check = [tv](int x) { return x <= tv; };
68-
size_t ans = a.rend() - std::find_if(a.rend() - l - 1, a.rend(), check) - 1;
69-
size_t out = tree.find_last_left(l, check);
70-
EXPECT_EQ(out, ans);
71-
}
72-
}

0 commit comments

Comments
 (0)