Skip to content

Commit 44b3805

Browse files
qryxipqryxip
authored
Merge pull request #101 from TonalidadeHidrica/range-bounds
2 parents c48fcaa + 686372b commit 44b3805

7 files changed

+152
-28
lines changed

‎examples/library-checker-static-range-sum.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,6 @@ fn main() {
2020
fenwick.add(i, a);
2121
}
2222
for (l, r) in lrs {
23-
println!("{}", fenwick.sum(l, r));
23+
println!("{}", fenwick.sum(l..r));
2424
}
2525
}

‎examples/practice2_j_segment_tree.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,9 @@ fn main() {
2020
segtree.set(x, v);
2121
}
2222
2 => {
23-
let l = input.next().unwrap().parse().unwrap();
23+
let l: usize = input.next().unwrap().parse().unwrap();
2424
let r: usize = input.next().unwrap().parse().unwrap();
25-
println!("{}", segtree.prod(l, r + 1));
25+
println!("{}", segtree.prod(l..=r));
2626
}
2727
3 => {
2828
let x = input.next().unwrap().parse().unwrap();

‎examples/practice2_k_range_affine_range_sum.rs

+5-5
Original file line numberDiff line numberDiff line change
@@ -50,16 +50,16 @@ fn main() {
5050
for _ in 0..q {
5151
match input.next().unwrap().parse().unwrap() {
5252
0 => {
53-
let l = input.next().unwrap().parse().unwrap();
53+
let l: usize = input.next().unwrap().parse().unwrap();
5454
let r = input.next().unwrap().parse().unwrap();
5555
let b = input.next().unwrap().parse().unwrap();
5656
let c = input.next().unwrap().parse().unwrap();
57-
segtree.apply_range(l, r, (b, c));
57+
segtree.apply_range(l..r, (b, c));
5858
}
5959
1 => {
60-
let l = input.next().unwrap().parse().unwrap();
61-
let r = input.next().unwrap().parse().unwrap();
62-
println!("{}", segtree.prod(l, r).0);
60+
let l: usize = input.next().unwrap().parse().unwrap();
61+
let r: usize = input.next().unwrap().parse().unwrap();
62+
println!("{}", segtree.prod(l..r).0);
6363
}
6464
_ => {}
6565
}

‎examples/practice2_l_lazy_segment_tree.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -55,8 +55,8 @@ fn main() {
5555
let l = input.next().unwrap().parse().unwrap();
5656
let r: usize = input.next().unwrap().parse().unwrap();
5757
match t {
58-
1 => segtree.apply_range(l, r + 1, true),
59-
2 => println!("{}", segtree.prod(l, r + 1).2),
58+
1 => segtree.apply_range(l..=r, true),
59+
2 => println!("{}", segtree.prod(l..=r).2),
6060
_ => {}
6161
}
6262
}

‎src/fenwicktree.rs

+25-4
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
use std::ops::{Bound, RangeBounds};
2+
13
// Reference: https://en.wikipedia.org/wiki/Fenwick_tree
24
pub struct FenwickTree<T> {
35
n: usize,
@@ -34,17 +36,29 @@ impl<T: Clone + std::ops::AddAssign<T>> FenwickTree<T> {
3436
}
3537
}
3638
/// Returns data[l] + ... + data[r - 1].
37-
pub fn sum(&self, l: usize, r: usize) -> T
39+
pub fn sum<R>(&self, range: R) -> T
3840
where
3941
T: std::ops::Sub<Output = T>,
42+
R: RangeBounds<usize>,
4043
{
44+
let r = match range.end_bound() {
45+
Bound::Included(r) => r + 1,
46+
Bound::Excluded(r) => *r,
47+
Bound::Unbounded => self.n,
48+
};
49+
let l = match range.start_bound() {
50+
Bound::Included(l) => *l,
51+
Bound::Excluded(l) => l + 1,
52+
Bound::Unbounded => return self.accum(r),
53+
};
4154
self.accum(r) - self.accum(l)
4255
}
4356
}
4457

4558
#[cfg(test)]
4659
mod tests {
4760
use super::*;
61+
use std::ops::Bound::*;
4862

4963
#[test]
5064
fn fenwick_tree_works() {
@@ -53,8 +67,15 @@ mod tests {
5367
for i in 0..5 {
5468
bit.add(i, i as i64 + 1);
5569
}
56-
assert_eq!(bit.sum(0, 5), 15);
57-
assert_eq!(bit.sum(0, 4), 10);
58-
assert_eq!(bit.sum(1, 3), 5);
70+
assert_eq!(bit.sum(0..5), 15);
71+
assert_eq!(bit.sum(0..4), 10);
72+
assert_eq!(bit.sum(1..3), 5);
73+
74+
assert_eq!(bit.sum(..), 15);
75+
assert_eq!(bit.sum(..2), 3);
76+
assert_eq!(bit.sum(..=2), 6);
77+
assert_eq!(bit.sum(1..), 14);
78+
assert_eq!(bit.sum(1..=3), 9);
79+
assert_eq!(bit.sum((Excluded(0), Included(2))), 5);
5980
}
6081
}

‎src/lazysegtree.rs

+71-8
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,27 @@ impl<F: MapMonoid> LazySegtree<F> {
7373
self.d[p].clone()
7474
}
7575

76-
pub fn prod(&mut self, mut l: usize, mut r: usize) -> <F::M as Monoid>::S {
76+
pub fn prod<R>(&mut self, range: R) -> <F::M as Monoid>::S
77+
where
78+
R: RangeBounds<usize>,
79+
{
80+
// Trivial optimization
81+
if range.start_bound() == Bound::Unbounded && range.end_bound() == Bound::Unbounded {
82+
return self.all_prod();
83+
}
84+
85+
let mut r = match range.end_bound() {
86+
Bound::Included(r) => r + 1,
87+
Bound::Excluded(r) => *r,
88+
Bound::Unbounded => self.n,
89+
};
90+
let mut l = match range.start_bound() {
91+
Bound::Included(l) => *l,
92+
Bound::Excluded(l) => l + 1,
93+
// TODO: There are another way of optimizing [0..r)
94+
Bound::Unbounded => 0,
95+
};
96+
7797
assert!(l <= r && r <= self.n);
7898
if l == r {
7999
return F::identity_element();
@@ -124,7 +144,22 @@ impl<F: MapMonoid> LazySegtree<F> {
124144
self.update(p >> i);
125145
}
126146
}
127-
pub fn apply_range(&mut self, mut l: usize, mut r: usize, f: F::F) {
147+
pub fn apply_range<R>(&mut self, range: R, f: F::F)
148+
where
149+
R: RangeBounds<usize>,
150+
{
151+
let mut r = match range.end_bound() {
152+
Bound::Included(r) => r + 1,
153+
Bound::Excluded(r) => *r,
154+
Bound::Unbounded => self.n,
155+
};
156+
let mut l = match range.start_bound() {
157+
Bound::Included(l) => *l,
158+
Bound::Excluded(l) => l + 1,
159+
// TODO: There are another way of optimizing [0..r)
160+
Bound::Unbounded => 0,
161+
};
162+
128163
assert!(l <= r && r <= self.n);
129164
if l == r {
130165
return;
@@ -287,7 +322,10 @@ where
287322
}
288323

289324
// TODO is it useful?
290-
use std::fmt::{Debug, Error, Formatter, Write};
325+
use std::{
326+
fmt::{Debug, Error, Formatter, Write},
327+
ops::{Bound, RangeBounds},
328+
};
291329
impl<F> Debug for LazySegtree<F>
292330
where
293331
F: MapMonoid,
@@ -314,6 +352,8 @@ where
314352

315353
#[cfg(test)]
316354
mod tests {
355+
use std::ops::{Bound::*, RangeBounds};
356+
317357
use crate::{LazySegtree, MapMonoid, Max};
318358

319359
struct MaxAdd;
@@ -361,9 +401,13 @@ mod tests {
361401
internal[6] = 0;
362402
check_segtree(&internal, &mut segtree);
363403

364-
segtree.apply_range(3, 8, 2);
404+
segtree.apply_range(3..8, 2);
365405
internal[3..8].iter_mut().for_each(|e| *e += 2);
366406
check_segtree(&internal, &mut segtree);
407+
408+
segtree.apply_range(2..=5, 7);
409+
internal[2..=5].iter_mut().for_each(|e| *e += 7);
410+
check_segtree(&internal, &mut segtree);
367411
}
368412

369413
//noinspection DuplicatedCode
@@ -373,12 +417,20 @@ mod tests {
373417
for i in 0..n {
374418
assert_eq!(segtree.get(i), base[i]);
375419
}
420+
421+
check(base, segtree, ..);
376422
for i in 0..=n {
423+
check(base, segtree, ..i);
424+
check(base, segtree, i..);
425+
if i < n {
426+
check(base, segtree, ..=i);
427+
}
377428
for j in i..=n {
378-
assert_eq!(
379-
segtree.prod(i, j),
380-
base[i..j].iter().max().copied().unwrap_or(i32::min_value())
381-
);
429+
check(base, segtree, i..j);
430+
if j < n {
431+
check(base, segtree, i..=j);
432+
check(base, segtree, (Excluded(i), Included(j)));
433+
}
382434
}
383435
}
384436
assert_eq!(
@@ -413,4 +465,15 @@ mod tests {
413465
}
414466
}
415467
}
468+
469+
fn check(base: &[i32], segtree: &mut LazySegtree<MaxAdd>, range: impl RangeBounds<usize>) {
470+
let expected = base
471+
.iter()
472+
.enumerate()
473+
.filter_map(|(i, a)| Some(a).filter(|_| range.contains(&i)))
474+
.max()
475+
.copied()
476+
.unwrap_or(i32::min_value());
477+
assert_eq!(segtree.prod(range), expected);
478+
}
416479
}

‎src/segtree.rs

+46-6
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use crate::internal_type_traits::{BoundedAbove, BoundedBelow, One, Zero};
33
use std::cmp::{max, min};
44
use std::convert::Infallible;
55
use std::marker::PhantomData;
6-
use std::ops::{Add, Mul};
6+
use std::ops::{Add, Bound, Mul, RangeBounds};
77

88
// TODO Should I split monoid-related traits to another module?
99
pub trait Monoid {
@@ -107,7 +107,27 @@ impl<M: Monoid> Segtree<M> {
107107
self.d[p + self.size].clone()
108108
}
109109

110-
pub fn prod(&self, mut l: usize, mut r: usize) -> M::S {
110+
pub fn prod<R>(&self, range: R) -> M::S
111+
where
112+
R: RangeBounds<usize>,
113+
{
114+
// Trivial optimization
115+
if range.start_bound() == Bound::Unbounded && range.end_bound() == Bound::Unbounded {
116+
return self.all_prod();
117+
}
118+
119+
let mut r = match range.end_bound() {
120+
Bound::Included(r) => r + 1,
121+
Bound::Excluded(r) => *r,
122+
Bound::Unbounded => self.n,
123+
};
124+
let mut l = match range.start_bound() {
125+
Bound::Included(l) => *l,
126+
Bound::Excluded(l) => l + 1,
127+
// TODO: There are another way of optimizing [0..r)
128+
Bound::Unbounded => 0,
129+
};
130+
111131
assert!(l <= r && r <= self.n);
112132
let mut sml = M::identity();
113133
let mut smr = M::identity();
@@ -240,6 +260,7 @@ where
240260
mod tests {
241261
use crate::segtree::Max;
242262
use crate::Segtree;
263+
use std::ops::{Bound::*, RangeBounds};
243264

244265
#[test]
245266
fn test_max_segtree() {
@@ -272,12 +293,20 @@ mod tests {
272293
for i in 0..n {
273294
assert_eq!(segtree.get(i), base[i]);
274295
}
296+
297+
check(base, segtree, ..);
275298
for i in 0..=n {
299+
check(base, segtree, ..i);
300+
check(base, segtree, i..);
301+
if i < n {
302+
check(base, segtree, ..=i);
303+
}
276304
for j in i..=n {
277-
assert_eq!(
278-
segtree.prod(i, j),
279-
base[i..j].iter().max().copied().unwrap_or(i32::min_value())
280-
);
305+
check(base, segtree, i..j);
306+
if j < n {
307+
check(base, segtree, i..=j);
308+
check(base, segtree, (Excluded(i), Included(j)));
309+
}
281310
}
282311
}
283312
assert_eq!(
@@ -312,4 +341,15 @@ mod tests {
312341
}
313342
}
314343
}
344+
345+
fn check(base: &[i32], segtree: &Segtree<Max<i32>>, range: impl RangeBounds<usize>) {
346+
let expected = base
347+
.iter()
348+
.enumerate()
349+
.filter_map(|(i, a)| Some(a).filter(|_| range.contains(&i)))
350+
.max()
351+
.copied()
352+
.unwrap_or(i32::min_value());
353+
assert_eq!(segtree.prod(range), expected);
354+
}
315355
}

0 commit comments

Comments
 (0)