1use super::Monoid;
2use std::{
3 fmt::{self, Debug},
4 marker::PhantomData,
5 ops::{Bound, RangeBounds},
6};
7
8pub struct CompressedSegmentTree<M, X, Inner>
9where
10 M: Monoid,
11{
12 compress: Vec<X>,
13 segs: Vec<Inner>,
14 _marker: PhantomData<fn() -> M>,
15}
16
17impl<M, X, Inner> Debug for CompressedSegmentTree<M, X, Inner>
18where
19 M: Monoid,
20 X: Debug,
21 Inner: Debug,
22{
23 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
24 f.debug_struct("CompressedSegmentTree")
25 .field("compress", &self.compress)
26 .field("segs", &self.segs)
27 .finish()
28 }
29}
30
31impl<M, X, Inner> Clone for CompressedSegmentTree<M, X, Inner>
32where
33 M: Monoid,
34 X: Clone,
35 Inner: Clone,
36{
37 fn clone(&self) -> Self {
38 Self {
39 compress: self.compress.clone(),
40 segs: self.segs.clone(),
41 _marker: self._marker,
42 }
43 }
44}
45
46impl<M, X, Inner> Default for CompressedSegmentTree<M, X, Inner>
47where
48 M: Monoid,
49{
50 fn default() -> Self {
51 Self {
52 compress: Default::default(),
53 segs: Default::default(),
54 _marker: Default::default(),
55 }
56 }
57}
58
59#[repr(transparent)]
60pub struct Tag<M>(M::T)
61where
62 M: Monoid;
63
64impl<M> Debug for Tag<M>
65where
66 M: Monoid<T: Debug>,
67{
68 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
69 self.0.fmt(f)
70 }
71}
72
73impl<M> Clone for Tag<M>
74where
75 M: Monoid,
76{
77 fn clone(&self) -> Self {
78 Self(self.0.clone())
79 }
80}
81
82impl<M, X> CompressedSegmentTree<M, X, Tag<M>>
83where
84 M: Monoid,
85 X: Clone + Ord,
86{
87 fn merge_1d_coordinates(left: &Self, right: &Self) -> Self {
88 let left = &left.compress;
89 let right = &right.compress;
90 let mut compress = Vec::with_capacity(left.len() + right.len());
91 let mut l = 0;
92 let mut r = 0;
93 while l < left.len() && r < right.len() {
94 match left[l].cmp(&right[r]) {
95 std::cmp::Ordering::Less => {
96 compress.push(left[l].clone());
97 l += 1;
98 }
99 std::cmp::Ordering::Equal => {
100 compress.push(left[l].clone());
101 l += 1;
102 r += 1;
103 }
104 std::cmp::Ordering::Greater => {
105 compress.push(right[r].clone());
106 r += 1;
107 }
108 }
109 }
110 compress.extend_from_slice(&left[l..]);
111 compress.extend_from_slice(&right[r..]);
112 let n = compress.len();
113 Self {
114 compress,
115 segs: vec![Tag(M::unit()); n * 2],
116 _marker: PhantomData,
117 }
118 }
119}
120
121trait MergeCoordinates<M>: Clone
122where
123 M: Monoid,
124{
125 fn empty() -> Self;
126 fn merge_coordinates(left: &Self, right: &Self) -> Self;
127}
128
129impl<M> MergeCoordinates<M> for Tag<M>
130where
131 M: Monoid,
132{
133 fn empty() -> Self {
134 Self(M::unit())
135 }
136
137 fn merge_coordinates(_: &Self, _: &Self) -> Self {
138 Self::empty()
139 }
140}
141
142impl<M, X, Inner> MergeCoordinates<M> for CompressedSegmentTree<M, X, Inner>
143where
144 M: Monoid,
145 X: Clone + Ord,
146 Inner: MergeCoordinates<M>,
147{
148 fn empty() -> Self {
149 Self::default()
150 }
151
152 fn merge_coordinates(left: &Self, right: &Self) -> Self {
153 let left_values = &left.compress;
154 let right_values = &right.compress;
155 let mut compress = Vec::with_capacity(left_values.len() + right_values.len());
156 let mut leaves = Vec::with_capacity(compress.capacity());
157 let mut l = 0;
158 let mut r = 0;
159 while l < left_values.len() && r < right_values.len() {
160 match left_values[l].cmp(&right_values[r]) {
161 std::cmp::Ordering::Less => {
162 compress.push(left_values[l].clone());
163 leaves.push(left.segs[left_values.len() + l].clone());
164 l += 1;
165 }
166 std::cmp::Ordering::Equal => {
167 compress.push(left_values[l].clone());
168 leaves.push(Inner::merge_coordinates(
169 &left.segs[left_values.len() + l],
170 &right.segs[right_values.len() + r],
171 ));
172 l += 1;
173 r += 1;
174 }
175 std::cmp::Ordering::Greater => {
176 compress.push(right_values[r].clone());
177 leaves.push(right.segs[right_values.len() + r].clone());
178 r += 1;
179 }
180 }
181 }
182 for (coordinate, leaf) in left_values[l..]
183 .iter()
184 .zip(&left.segs[left_values.len() + l..])
185 {
186 compress.push(coordinate.clone());
187 leaves.push(leaf.clone());
188 }
189 for (coordinate, leaf) in right_values[r..]
190 .iter()
191 .zip(&right.segs[right_values.len() + r..])
192 {
193 compress.push(coordinate.clone());
194 leaves.push(leaf.clone());
195 }
196 let n = compress.len();
197 let mut segs = vec![Inner::empty(); n];
198 segs.extend(leaves);
199 for i in (1..n).rev() {
200 segs[i] = Inner::merge_coordinates(&segs[i * 2], &segs[i * 2 + 1]);
201 }
202 Self {
203 compress,
204 segs,
205 _marker: PhantomData,
206 }
207 }
208}
209
210macro_rules! impl_compressed_segment_tree {
211 (@tuple ($($l:tt)*) ($($r:tt)*) $T:ident) => {
212 ($($l)* $T $($r)*,)
213 };
214 (@tuple ($($l:tt)*) ($($r:tt)*) $T:ident $($Rest:ident)+) => {
215 ($($l)* $T $($r)*, impl_compressed_segment_tree!(@tuple ($($l)*) ($($r)*) $($Rest)+))
216 };
217 (@cst $M:ident) => {
218 Tag<$M>
219 };
220 (@cst $M:ident $T:ident $($Rest:ident)*) => {
221 CompressedSegmentTree<$M, $T, impl_compressed_segment_tree!(@cst $M $($Rest)*)>
222 };
223 (@from_iter $M:ident $points:ident $T:ident) => {{
224 let mut compress: Vec<_> = $points.into_iter().map(|t| t.0.clone()).collect();
225 compress.sort_unstable();
226 compress.dedup();
227 let n = compress.len();
228 Self {
229 compress,
230 segs: vec![Tag(M::unit()); n * 2],
231 _marker: PhantomData,
232 }
233 }};
234 (@from_iter $M:ident $points:ident $T:ident $U:ident) => {{
235 let mut points: Vec<_> = $points.into_iter().collect();
236 points.sort_unstable_by(|left, right| left.0.cmp(&right.0));
237 let mut compress = Vec::new();
238 let mut offsets = vec![0];
239 let mut start = 0;
240 while start < points.len() {
241 let mut end = start + 1;
242 while end < points.len() && points[end].0 == points[start].0 {
243 end += 1;
244 }
245 compress.push(points[start].0.clone());
246 offsets.push(end);
247 start = end;
248 }
249 let n = compress.len();
250 let mut segs: Vec<impl_compressed_segment_tree!(@cst $M $U)> =
251 vec![Default::default(); n];
252 for i in 0..n {
253 segs.push(<impl_compressed_segment_tree!(@cst $M $U)>::from_iter(
254 points[offsets[i]..offsets[i + 1]]
255 .iter()
256 .map(|point| &point.1),
257 ));
258 }
259 for i in (1..n).rev() {
260 segs[i] = CompressedSegmentTree::merge_1d_coordinates(
261 &segs[i * 2],
262 &segs[i * 2 + 1],
263 );
264 }
265 Self {
266 compress,
267 segs,
268 _marker: PhantomData,
269 }
270 }};
271 (@from_iter $M:ident $points:ident $T:ident $U:ident $($Rest:ident)*) => {{
272 let mut points: Vec<_> = $points.into_iter().collect();
273 points.sort_unstable_by(|left, right| left.0.cmp(&right.0));
274 let mut compress = Vec::new();
275 let mut offsets = vec![0];
276 let mut start = 0;
277 while start < points.len() {
278 let mut end = start + 1;
279 while end < points.len() && points[end].0 == points[start].0 {
280 end += 1;
281 }
282 compress.push(points[start].0.clone());
283 offsets.push(end);
284 start = end;
285 }
286 let n = compress.len();
287 let mut segs: Vec<impl_compressed_segment_tree!(@cst $M $U $($Rest)*)> =
288 vec![Default::default(); n];
289 for i in 0..n {
290 segs.push(<impl_compressed_segment_tree!(
291 @cst $M $U $($Rest)*
292 )>::from_iter(
293 points[offsets[i]..offsets[i + 1]]
294 .iter()
295 .map(|point| &point.1),
296 ));
297 }
298 for i in (1..n).rev() {
299 segs[i] = MergeCoordinates::merge_coordinates(&segs[i * 2], &segs[i * 2 + 1]);
300 }
301 Self {
302 compress,
303 segs,
304 _marker: PhantomData,
305 }
306 }};
307 (@fold $e:expr, $rng:ident $T:ident) => {
308 $e.0
309 };
310 (@fold $e:expr, $rng:ident $T:ident $($Rest:ident)+) => {
311 $e.fold(&$rng.1)
312 };
313 (@update $e:expr, $M:ident $key:ident $x:ident $T:ident) => {
314 $M::operate_assign(&mut $e.0, $x);
315 };
316 (@update $e:expr, $M:ident $key:ident $x:ident $T:ident $($Rest:ident)+) => {
317 $e.update(&$key.1, $x);
318 };
319 (@partition_value $e:expr) => {
320 $e.0
321 };
322 (@partition_value $e:expr, $inner_ranges:ident) => {
323 $e.fold($inner_ranges)
324 };
325 (@partition_perfect_body $this:tt, $pos:expr, $acc:expr, $pred:ident $(, $inner_ranges:ident)?) => {{
326 let n = $this.compress.len();
327 let mut pos = $pos;
328 let mut acc = $acc;
329 while pos < n {
330 pos <<= 1;
331 let nacc = M::operate(
332 &acc,
333 &impl_compressed_segment_tree!(
334 @partition_value $this.segs[pos] $(, $inner_ranges)?
335 ),
336 );
337 if $pred(&nacc) {
338 acc = nacc;
339 pos += 1;
340 }
341 }
342 (pos - n, acc)
343 }};
344 (@rpartition_perfect_body $this:tt, $pos:expr, $acc:expr, $pred:ident $(, $inner_ranges:ident)?) => {{
345 let n = $this.compress.len();
346 let mut pos = $pos;
347 let mut acc = $acc;
348 while pos < n {
349 pos = pos * 2 + 1;
350 let nacc = M::operate(
351 &impl_compressed_segment_tree!(
352 @partition_value $this.segs[pos] $(, $inner_ranges)?
353 ),
354 &acc,
355 );
356 if $pred(&nacc) {
357 acc = nacc;
358 pos -= 1;
359 }
360 }
361 (pos - n + 1, acc)
362 }};
363 (@partition_body $this:tt, $left:ident, $pred:ident $(, $inner_ranges:ident)?) => {{
364 let n = $this.compress.len();
365 let mut l = $this.compress.partition_point(|x| x < $left) + n;
366 let r = 2 * n;
367 let mut k = 0usize;
368 let mut acc = M::unit();
369 while l < r >> k {
370 if l & 1 != 0 {
371 let nacc = M::operate(
372 &acc,
373 &impl_compressed_segment_tree!(
374 @partition_value $this.segs[l] $(, $inner_ranges)?
375 ),
376 );
377 if !$pred(&nacc) {
378 let (pos, acc) = impl_compressed_segment_tree!(
379 @partition_perfect_body $this, l, acc, $pred $(, $inner_ranges)?
380 );
381 return ($this.compress.get(pos), acc);
382 }
383 acc = nacc;
384 l += 1;
385 }
386 l >>= 1;
387 k += 1;
388 }
389 for k in (0..k).rev() {
390 let r = r >> k;
391 if r & 1 != 0 {
392 let nacc = M::operate(
393 &acc,
394 &impl_compressed_segment_tree!(
395 @partition_value $this.segs[r - 1] $(, $inner_ranges)?
396 ),
397 );
398 if !$pred(&nacc) {
399 let (pos, acc) = impl_compressed_segment_tree!(
400 @partition_perfect_body $this, r - 1, acc, $pred $(, $inner_ranges)?
401 );
402 return ($this.compress.get(pos), acc);
403 }
404 acc = nacc;
405 }
406 }
407 ($this.compress.get(n), acc)
408 }};
409 (@rpartition_body $this:tt, $right:ident, $pred:ident $(, $inner_ranges:ident)?) => {{
410 let n = $this.compress.len();
411 let mut l = n;
412 let mut r = $this.compress.partition_point(|x| x < $right) + n;
413 let mut c = 0usize;
414 let mut k = 0usize;
415 let mut acc = M::unit();
416 while l >> k < r {
417 c <<= 1;
418 if l & (1 << k) != 0 {
419 l += 1 << k;
420 c += 1;
421 }
422 if r & 1 != 0 {
423 r -= 1;
424 let nacc = M::operate(
425 &impl_compressed_segment_tree!(
426 @partition_value $this.segs[r] $(, $inner_ranges)?
427 ),
428 &acc,
429 );
430 if !$pred(&nacc) {
431 let (pos, acc) = impl_compressed_segment_tree!(
432 @rpartition_perfect_body $this, r, acc, $pred $(, $inner_ranges)?
433 );
434 return ($this.compress.get(pos), acc);
435 }
436 acc = nacc;
437 }
438 r >>= 1;
439 k += 1;
440 }
441 for k in (0..k).rev() {
442 if c & 1 != 0 {
443 l -= 1 << k;
444 let l = l >> k;
445 let nacc = M::operate(
446 &impl_compressed_segment_tree!(
447 @partition_value $this.segs[l] $(, $inner_ranges)?
448 ),
449 &acc,
450 );
451 if !$pred(&nacc) {
452 let (pos, acc) = impl_compressed_segment_tree!(
453 @rpartition_perfect_body $this, l, acc, $pred $(, $inner_ranges)?
454 );
455 return ($this.compress.get(pos), acc);
456 }
457 acc = nacc;
458 }
459 c >>= 1;
460 }
461 ($this.compress.first(), acc)
462 }};
463 (@partition_methods $T:ident, $Q:ident) => {
464 pub fn partition_point_acc<P>(
465 &self,
466 left: &$T,
467 mut pred: P,
468 ) -> (Option<&$T>, M::T)
469 where
470 P: FnMut(&M::T) -> bool,
471 {
472 impl_compressed_segment_tree!(@partition_body self, left, pred)
473 }
474 pub fn rpartition_point_acc<P>(
475 &self,
476 right: &$T,
477 mut pred: P,
478 ) -> (Option<&$T>, M::T)
479 where
480 P: FnMut(&M::T) -> bool,
481 {
482 impl_compressed_segment_tree!(@rpartition_body self, right, pred)
483 }
484 };
485 (@partition_methods $T:ident $($RestT:ident)+, $Q:ident $($RestQ:ident)+) => {
486 pub fn partition_point_acc<P, $($RestQ,)*>(
487 &self,
488 left: &$T,
489 inner_ranges: &impl_compressed_segment_tree!(@tuple () () $($RestQ)*),
490 mut pred: P,
491 ) -> (Option<&$T>, M::T)
492 where
493 P: FnMut(&M::T) -> bool,
494 $($RestQ: RangeBounds<$RestT>,)*
495 {
496 impl_compressed_segment_tree!(
497 @partition_body self, left, pred, inner_ranges
498 )
499 }
500 pub fn rpartition_point_acc<P, $($RestQ,)*>(
501 &self,
502 right: &$T,
503 inner_ranges: &impl_compressed_segment_tree!(@tuple () () $($RestQ)*),
504 mut pred: P,
505 ) -> (Option<&$T>, M::T)
506 where
507 P: FnMut(&M::T) -> bool,
508 $($RestQ: RangeBounds<$RestT>,)*
509 {
510 impl_compressed_segment_tree!(
511 @rpartition_body self, right, pred, inner_ranges
512 )
513 }
514 };
515 (@impl $C:ident $($T:ident)*, $($Q:ident)*) => {
516 impl<M, $($T,)*> impl_compressed_segment_tree!(@cst M $($T)*)
517 where
518 M: Monoid,
519 $($T: Clone + Ord,)*
520 {
521 pub fn new(points: &[impl_compressed_segment_tree!(@tuple () () $($T)*)]) -> Self {
522 Self::from_iter(points)
523 }
524 fn from_iter<'a, Iter>(points: Iter) -> Self
525 where
526 $($T: 'a,)*
527 Iter: IntoIterator<Item = &'a impl_compressed_segment_tree!(@tuple () () $($T)*)> + Clone,
528 {
529 impl_compressed_segment_tree!(@from_iter M points $($T)*)
530 }
531 pub fn fold<$($Q,)*>(&self, range: &impl_compressed_segment_tree!(@tuple () () $($Q)*)) -> M::T
532 where
533 $($Q: RangeBounds<$T>,)*
534 {
535 let mut l = match range.0.start_bound() {
536 Bound::Included(index) => self.compress.partition_point(|x| x < index),
537 Bound::Excluded(index) => self.compress.partition_point(|x| x <= index),
538 Bound::Unbounded => 0,
539 } + self.compress.len();
540 let mut r = match range.0.end_bound() {
541 Bound::Included(index) => self.compress.partition_point(|x| x <= index),
542 Bound::Excluded(index) => self.compress.partition_point(|x| x < index),
543 Bound::Unbounded => self.compress.len(),
544 } + self.compress.len();
545 let mut x = M::unit();
546 while l < r {
547 if l & 1 != 0 {
548 x = M::operate(&x, &impl_compressed_segment_tree!(@fold self.segs[l], range $($T)*));
549 l += 1;
550 }
551 if r & 1 != 0 {
552 r -= 1;
553 x = M::operate(&impl_compressed_segment_tree!(@fold self.segs[r], range $($T)*), &x);
554 }
555 l /= 2;
556 r /= 2;
557 }
558 x
559 }
560 pub fn update(&mut self, key: &impl_compressed_segment_tree!(@tuple () () $($T)*), x: &M::T) {
561 let mut i = self.compress.binary_search(&key.0).expect("not exist key") + self.compress.len();
562 while i > 0 {
563 impl_compressed_segment_tree!(@update self.segs[i], M key x $($T)*);
564 i /= 2;
565 }
566 }
567 impl_compressed_segment_tree!(
568 @partition_methods $($T)*, $($Q)*
569 );
570 }
571 pub type $C<M, $($T),*> = impl_compressed_segment_tree!(@cst M $($T)*);
572 };
573 (@inner [$C:ident][$($T:ident)*][$($Q:ident)*][]) => {
574 impl_compressed_segment_tree!(@impl $C $($T)*, $($Q)*);
575 };
576 (@inner [$C:ident][$($T:ident)*][$($Q:ident)*][$D:ident $U:ident $R:ident $($Rest:ident)*]) => {
577 impl_compressed_segment_tree!(@impl $C $($T)*, $($Q)*);
578 impl_compressed_segment_tree!(@inner [$D][$($T)* $U][$($Q)* $R][$($Rest)*]);
579 };
580 ($C:ident $T:ident $Q:ident $($Rest:ident)* $(;$($t:tt)*)?) => {
581 impl_compressed_segment_tree!(@inner [$C][$T][$Q][$($Rest)*]);
582 };
583}
584
585impl_compressed_segment_tree!(
586 CompressedSegmentTree1d T1 Q1
587 CompressedSegmentTree2d T2 Q2
588 CompressedSegmentTree3d T3 Q3
589 CompressedSegmentTree4d T4 Q4;
590 CompressedSegmentTree5d T5 Q5
591 CompressedSegmentTree6d T6 Q6
592 CompressedSegmentTree7d T7 Q7
593 CompressedSegmentTree8d T8 Q8
594 CompressedSegmentTree9d T9 Q9
595);
596
597#[cfg(test)]
598mod tests {
599 use super::*;
600 use crate::{
601 algebra::AdditiveOperation,
602 tools::{RandRange as RR, Xorshift},
603 };
604 use std::{collections::HashMap, ops::Range};
605
606 #[test]
607 fn test_seg1d_partition_point_acc() {
608 let mut rng = Xorshift::default();
609 const N: usize = 100;
610 const Q: usize = 5000;
611 const A: Range<u64> = 0..1_000;
612 let mut points: Vec<_> = rng.random_iter(A).take(N).map(|x| (x,)).collect();
613 points.sort();
614 points.dedup();
615 let mut values: HashMap<_, _> = points.iter().map(|p| (p.0, 0u64)).collect();
616 let mut seg = CompressedSegmentTree1d::<AdditiveOperation<u64>, _>::new(&points);
617 for _ in 0..Q {
618 let p = &points[rng.random(0..points.len())];
619 let x = rng.random(A);
620 *values.get_mut(&p.0).unwrap() += x;
621 seg.update(p, &x);
622
623 let left = rng.random(A);
624 let target = rng.random(1..A.end * Q as u64);
625 let mut expected_acc = 0;
626 let compress = &seg.compress;
627 let mut expected_pos = compress.partition_point(|x| x < &left);
628 while expected_pos < compress.len() {
629 let nacc = expected_acc + values[&compress[expected_pos]];
630 if nacc < target {
631 expected_acc = nacc;
632 expected_pos += 1;
633 } else {
634 break;
635 }
636 }
637 let (pos, acc) = seg.partition_point_acc(&left, |&acc| acc < target);
638 assert_eq!((pos, acc), (compress.get(expected_pos), expected_acc));
639
640 let right = rng.random(A);
641 let target = rng.random(1..A.end * Q as u64);
642 let mut expected_acc = 0;
643 let mut expected_pos = compress.partition_point(|x| x < &right);
644 while expected_pos > 0 {
645 let nacc = values[&compress[expected_pos - 1]] + expected_acc;
646 if nacc < target {
647 expected_acc = nacc;
648 expected_pos -= 1;
649 } else {
650 break;
651 }
652 }
653 let (pos, acc) = seg.rpartition_point_acc(&right, |&acc| acc < target);
654 assert_eq!((pos, acc), (compress.get(expected_pos), expected_acc));
655 }
656 }
657
658 #[test]
659 fn test_seg2d_and_4d() {
660 let mut rng = Xorshift::default();
661 for _ in 0..12 {
662 let radius = rng.rand(128) as i64 + 1;
663 let point_count = rng.rand(96) as usize + 1;
664 let registered: Vec<_> = rng
665 .random_iter((-radius..radius, (-radius..radius,)))
666 .take(point_count)
667 .collect();
668 let mut points = registered.clone();
669 points.sort_unstable();
670 points.dedup();
671 let mut values: HashMap<_, _> =
672 points.iter().copied().map(|point| (point, 0i64)).collect();
673 let mut seg = CompressedSegmentTree2d::<AdditiveOperation<i64>, _, _>::new(®istered);
674 let query_count = rng.rand(300) as usize + 300;
675 for _ in 0..query_count {
676 let point = &points[rng.rand(points.len() as u64) as usize];
677 let value = rng.rand((radius * 2) as u64) as i64 - radius;
678 *values.get_mut(point).unwrap() += value;
679 seg.update(point, &value);
680
681 let range = rng.random((RR::new(-radius..radius), (RR::new(-radius..radius),)));
682 let expected = values
683 .iter()
684 .filter_map(|((x, (y,)), value)| {
685 (RangeBounds::contains(&range.0, x) && RangeBounds::contains(&range.1.0, y))
686 .then_some(*value)
687 })
688 .sum();
689 assert_eq!(seg.fold(&range), expected);
690 }
691 }
692
693 const N: usize = 100;
694 const Q: usize = 5000;
695 const A: Range<i64> = -1_000..1_000;
696 let mut points: Vec<_> = rng.random_iter(((A), (A, (A, (A,))))).take(N).collect();
697 points.sort();
698 points.dedup();
699 let mut map: HashMap<_, _> = points.iter().map(|p| (p, 0i64)).collect();
700 let mut seg = CompressedSegmentTree4d::<AdditiveOperation<i64>, _, _, _, _>::new(&points);
701 for _ in 0..Q {
702 let p = &points[rng.random(0..points.len())];
703 let x = rng.random(A);
704 *map.get_mut(p).unwrap() += x;
705 seg.update(p, &x);
706
707 let range = rng.random((RR::new(A), (RR::new(A), (RR::new(A), (RR::new(A),)))));
708 let (r0, (r1, (r2, (r3,)))) = range;
709 let expected: i64 = map
710 .iter()
711 .filter_map(|((p0, (p1, (p2, (p3,)))), x)| {
712 if RangeBounds::contains(&r0, p0)
713 && RangeBounds::contains(&r1, p1)
714 && RangeBounds::contains(&r2, p2)
715 && RangeBounds::contains(&r3, p3)
716 {
717 Some(*x)
718 } else {
719 None
720 }
721 })
722 .sum();
723 let result = seg.fold(&range);
724 assert_eq!(expected, result);
725 }
726 }
727
728 #[test]
729 fn test_seg4d_partition_point_acc() {
730 let mut rng = Xorshift::default();
731 const N: usize = 100;
732 const Q: usize = 1000;
733 const A: Range<u64> = 0..1_000;
734 let mut points: Vec<_> = rng.random_iter(((A), (A, (A, (A,))))).take(N).collect();
735 points.sort();
736 points.dedup();
737 let mut values: HashMap<_, _> = points.iter().map(|p| (p, 0u64)).collect();
738 let mut seg = CompressedSegmentTree4d::<AdditiveOperation<u64>, _, _, _, _>::new(&points);
739 for _ in 0..Q {
740 let p = &points[rng.random(0..points.len())];
741 let x = rng.random(A);
742 *values.get_mut(p).unwrap() += x;
743 seg.update(p, &x);
744
745 let inner_ranges = rng.random((RR::new(A), (RR::new(A), (RR::new(A),))));
746 let (r1, (r2, (r3,))) = &inner_ranges;
747 let compress = &seg.compress;
748 let mut groups = vec![0; compress.len()];
749 for ((p0, (p1, (p2, (p3,)))), x) in &values {
750 if RangeBounds::contains(r1, p1)
751 && RangeBounds::contains(r2, p2)
752 && RangeBounds::contains(r3, p3)
753 {
754 groups[compress.binary_search(p0).unwrap()] += *x;
755 }
756 }
757
758 let left = rng.random(A);
759 let target = rng.random(1..A.end * Q as u64);
760 let mut expected_acc = 0;
761 let mut expected_pos = compress.partition_point(|x| x < &left);
762 while expected_pos < compress.len() {
763 let nacc = expected_acc + groups[expected_pos];
764 if nacc < target {
765 expected_acc = nacc;
766 expected_pos += 1;
767 } else {
768 break;
769 }
770 }
771 let (pos, acc) = seg.partition_point_acc(&left, &inner_ranges, |&acc| acc < target);
772 assert_eq!((pos, acc), (compress.get(expected_pos), expected_acc));
773
774 let right = rng.random(A);
775 let target = rng.random(1..A.end * Q as u64);
776 let mut expected_acc = 0;
777 let mut expected_pos = compress.partition_point(|x| x < &right);
778 while expected_pos > 0 {
779 let nacc = groups[expected_pos - 1] + expected_acc;
780 if nacc < target {
781 expected_acc = nacc;
782 expected_pos -= 1;
783 } else {
784 break;
785 }
786 }
787 let (pos, acc) = seg.rpartition_point_acc(&right, &inner_ranges, |&acc| acc < target);
788 assert_eq!((pos, acc), (compress.get(expected_pos), expected_acc));
789 }
790 }
791}