1#[derive(Debug, Clone)]
37pub struct FenwickTree {
38 tree: Vec<u32>,
40 n: usize,
42}
43
44impl FenwickTree {
45 pub fn new(n: usize) -> Self {
47 Self {
48 tree: vec![0u32; n + 1],
49 n,
50 }
51 }
52
53 pub fn from_values(values: &[u32]) -> Self {
57 let n = values.len();
58 let mut tree = vec![0u32; n + 1];
59
60 for (i, &v) in values.iter().enumerate() {
62 tree[i + 1] = v;
63 }
64
65 for i in 1..=n {
67 let parent = i + lowbit(i);
68 if parent <= n {
69 tree[parent] = tree[parent].wrapping_add(tree[i]);
70 }
71 }
72
73 Self { tree, n }
74 }
75
76 #[inline]
78 pub fn len(&self) -> usize {
79 self.n
80 }
81
82 #[inline]
84 pub fn is_empty(&self) -> bool {
85 self.n == 0
86 }
87
88 pub fn update(&mut self, i: usize, delta: i32) {
93 self.update_u32(i, delta as u32);
98 }
99
100 pub fn set(&mut self, i: usize, value: u32) {
102 let current = self.get(i);
103 self.update_u32(i, value.wrapping_sub(current));
106 }
107
108 fn update_u32(&mut self, i: usize, delta: u32) {
110 assert!(i < self.n, "index {i} out of bounds (n={})", self.n);
111 let mut idx = i + 1; while idx <= self.n {
113 self.tree[idx] = self.tree[idx].wrapping_add(delta);
114 idx += lowbit(idx);
115 }
116 }
117
118 pub fn get(&self, i: usize) -> u32 {
122 if i == 0 {
123 self.prefix(0)
124 } else {
125 self.prefix(i).wrapping_sub(self.prefix(i - 1))
126 }
127 }
128
129 pub fn prefix(&self, i: usize) -> u32 {
134 assert!(i < self.n, "index {i} out of bounds (n={})", self.n);
135 let mut sum = 0u32;
136 let mut idx = i + 1; while idx > 0 {
138 sum = sum.wrapping_add(self.tree[idx]);
139 idx -= lowbit(idx);
140 }
141 sum
142 }
143
144 pub fn range(&self, left: usize, right: usize) -> u32 {
149 assert!(left <= right, "left {left} > right {right}");
150 if left == 0 {
151 self.prefix(right)
152 } else {
153 self.prefix(right).wrapping_sub(self.prefix(left - 1))
154 }
155 }
156
157 pub fn total(&self) -> u32 {
159 if self.n == 0 {
160 0
161 } else {
162 self.prefix(self.n - 1)
163 }
164 }
165
166 pub fn batch_update(&mut self, deltas: &[(usize, i32)]) {
171 for &(i, delta) in deltas {
172 self.update(i, delta);
173 }
174 }
175
176 pub fn rebuild(&mut self, values: &[u32]) {
183 assert_eq!(values.len(), self.n, "rebuild size mismatch");
184
185 self.tree.fill(0);
187
188 for (i, &v) in values.iter().enumerate() {
190 self.tree[i + 1] = v;
191 }
192
193 for i in 1..=self.n {
195 let parent = i + lowbit(i);
196 if parent <= self.n {
197 self.tree[parent] = self.tree[parent].wrapping_add(self.tree[i]);
198 }
199 }
200 }
201
202 #[must_use = "use the found index (if any)"]
208 pub fn find_prefix(&self, target: u32) -> Option<usize> {
209 if self.n == 0 {
210 return None;
211 }
212
213 let mut pos = 0usize;
214 let mut remaining = target;
215 let mut bit_mask = most_significant_bit(self.n);
216
217 while bit_mask > 0 {
218 let next = pos + bit_mask;
219 if next <= self.n && self.tree[next] <= remaining {
220 remaining -= self.tree[next];
221 pos = next;
222 }
223 bit_mask >>= 1;
224 }
225
226 if pos == 0 && self.tree.get(1).copied().unwrap_or(u32::MAX) > target {
227 None
228 } else {
229 Some(pos.saturating_sub(1)) }
231 }
232
233 pub fn resize(&mut self, new_n: usize) {
236 if new_n == self.n {
237 return;
238 }
239 let mut values: Vec<u32> = (0..self.n).map(|i| self.get(i)).collect();
241 values.resize(new_n, 0);
242 self.n = new_n;
243 self.tree = vec![0u32; new_n + 1];
244 self.rebuild(&values);
245 }
246}
247
248#[inline]
250fn lowbit(x: usize) -> usize {
251 x & x.wrapping_neg()
252}
253
254#[inline]
256fn most_significant_bit(n: usize) -> usize {
257 if n == 0 {
258 return 0;
259 }
260 1 << (usize::BITS - 1 - n.leading_zeros())
261}
262
263#[cfg(test)]
264mod tests {
265 use super::*;
266
267 #[test]
270 fn new_creates_zeroed_tree() {
271 let ft = FenwickTree::new(10);
272 assert_eq!(ft.len(), 10);
273 assert!(!ft.is_empty());
274 assert_eq!(ft.total(), 0);
275 }
276
277 #[test]
278 fn empty_tree() {
279 let ft = FenwickTree::new(0);
280 assert!(ft.is_empty());
281 assert_eq!(ft.total(), 0);
282 }
283
284 #[test]
285 fn from_values_matches_sequential() {
286 let values = vec![3, 1, 4, 1, 5, 9, 2, 6];
287 let ft = FenwickTree::from_values(&values);
288
289 assert_eq!(ft.prefix(0), 3);
291 assert_eq!(ft.prefix(1), 4);
292 assert_eq!(ft.prefix(2), 8);
293 assert_eq!(ft.prefix(7), 31);
294 assert_eq!(ft.total(), 31);
295 }
296
297 #[test]
300 fn update_and_query() {
301 let mut ft = FenwickTree::new(5);
302 ft.update(0, 10);
303 ft.update(2, 20);
304 ft.update(4, 30);
305
306 assert_eq!(ft.prefix(0), 10);
307 assert_eq!(ft.prefix(2), 30);
308 assert_eq!(ft.prefix(4), 60);
309 assert_eq!(ft.total(), 60);
310 }
311
312 #[test]
313 fn set_overwrites_value() {
314 let mut ft = FenwickTree::from_values(&[5, 10, 15]);
315 ft.set(1, 20);
316 assert_eq!(ft.get(0), 5);
317 assert_eq!(ft.get(1), 20);
318 assert_eq!(ft.get(2), 15);
319 assert_eq!(ft.total(), 40);
320 }
321
322 #[test]
323 fn get_retrieves_individual_values() {
324 let values = vec![7, 3, 8, 2, 6];
325 let ft = FenwickTree::from_values(&values);
326 for (i, &v) in values.iter().enumerate() {
327 assert_eq!(ft.get(i), v, "mismatch at index {i}");
328 }
329 }
330
331 #[test]
334 fn range_sum() {
335 let ft = FenwickTree::from_values(&[1, 2, 3, 4, 5]);
336 assert_eq!(ft.range(0, 4), 15);
337 assert_eq!(ft.range(1, 3), 9);
338 assert_eq!(ft.range(2, 2), 3);
339 assert_eq!(ft.range(0, 0), 1);
340 }
341
342 #[test]
345 fn unit_batch_update_equivalence() {
346 let values = vec![10, 20, 30, 40, 50];
347 let deltas = vec![(0, 5), (2, -3), (4, 10), (1, 7)];
348
349 let mut ft_seq = FenwickTree::from_values(&values);
351 for &(i, d) in &deltas {
352 ft_seq.update(i, d);
353 }
354
355 let mut ft_batch = FenwickTree::from_values(&values);
357 ft_batch.batch_update(&deltas);
358
359 for i in 0..5 {
361 assert_eq!(ft_seq.get(i), ft_batch.get(i), "mismatch at index {i}");
362 }
363 }
364
365 #[test]
368 fn rebuild_matches_from_values() {
369 let v1 = vec![1, 2, 3, 4, 5];
370 let v2 = vec![10, 20, 30, 40, 50];
371
372 let ft1 = FenwickTree::from_values(&v2);
373 let mut ft2 = FenwickTree::from_values(&v1);
374 ft2.rebuild(&v2);
375
376 for i in 0..5 {
377 assert_eq!(ft1.get(i), ft2.get(i));
378 }
379 }
380
381 #[test]
384 fn find_prefix_scroll_offset() {
385 let ft = FenwickTree::from_values(&[20, 30, 10, 40, 25]);
388
389 assert_eq!(ft.find_prefix(0), None);
393
394 assert_eq!(ft.find_prefix(20), Some(0));
396
397 assert_eq!(ft.find_prefix(50), Some(1));
399
400 assert_eq!(ft.find_prefix(99), Some(2));
402
403 assert_eq!(ft.find_prefix(125), Some(4));
405 }
406
407 #[test]
410 fn resize_grow_preserves() {
411 let mut ft = FenwickTree::from_values(&[1, 2, 3]);
412 ft.resize(5);
413 assert_eq!(ft.len(), 5);
414 assert_eq!(ft.get(0), 1);
415 assert_eq!(ft.get(1), 2);
416 assert_eq!(ft.get(2), 3);
417 assert_eq!(ft.get(3), 0);
418 assert_eq!(ft.get(4), 0);
419 }
420
421 #[test]
422 fn resize_shrink_drops() {
423 let mut ft = FenwickTree::from_values(&[1, 2, 3, 4, 5]);
424 ft.resize(3);
425 assert_eq!(ft.len(), 3);
426 assert_eq!(ft.total(), 6); }
428
429 #[test]
432 fn property_prefix_sum_correct() {
433 let mut seed: u64 = 0xCAFE_BABE_0000_0001;
435 let n = 100;
436 let mut naive = vec![0u32; n];
437 let mut ft = FenwickTree::new(n);
438
439 for _ in 0..500 {
440 seed = seed
441 .wrapping_mul(6364136223846793005)
442 .wrapping_add(1442695040888963407);
443 let idx = (seed >> 33) as usize % n;
444 seed = seed
445 .wrapping_mul(6364136223846793005)
446 .wrapping_add(1442695040888963407);
447 let delta = ((seed >> 33) as i32 % 100).abs();
448
449 naive[idx] = naive[idx].wrapping_add(delta as u32);
450 ft.update(idx, delta);
451 }
452
453 let mut naive_prefix = 0u32;
455 for (i, value) in naive.iter().enumerate() {
456 naive_prefix = naive_prefix.wrapping_add(*value);
457 assert_eq!(ft.prefix(i), naive_prefix, "prefix mismatch at index {i}");
458 }
459 }
460
461 #[test]
464 fn update_i32_min_does_not_panic() {
465 let mut ft = FenwickTree::from_values(&[0, 0, 0]);
467 ft.update(0, i32::MIN); assert_eq!(ft.get(0), i32::MIN as u32);
470 }
471
472 #[test]
473 fn set_large_u32_value() {
474 let mut ft = FenwickTree::from_values(&[0, 100, 200]);
476 ft.set(0, u32::MAX);
477 assert_eq!(ft.get(0), u32::MAX);
478 assert_eq!(ft.get(1), 100);
479 assert_eq!(ft.get(2), 200);
480 }
481
482 #[test]
485 fn perf_fenwick_hotpath() {
486 let n = 100_000;
487 let values: Vec<u32> = (0..n).map(|i| (i % 50 + 1) as u32).collect();
488
489 let start = std::time::Instant::now();
491 let mut ft = FenwickTree::from_values(&values);
492 let build_time = start.elapsed();
493
494 let start = std::time::Instant::now();
496 for i in 0..10_000 {
497 ft.update(i * 10, 5);
498 }
499 let update_time = start.elapsed();
500
501 let start = std::time::Instant::now();
503 let mut _sink = 0u32;
504 for i in 0..10_000 {
505 _sink = _sink.wrapping_add(ft.prefix(i * 10));
506 }
507 let query_time = start.elapsed();
508
509 let deltas: Vec<(usize, i32)> = (0..10_000).map(|i| (i * 10, 3)).collect();
511 let start = std::time::Instant::now();
512 ft.batch_update(&deltas);
513 let batch_time = start.elapsed();
514
515 eprintln!("=== Fenwick Tree Performance (n={n}) ===");
517 eprintln!("Build (from_values): {:?}", build_time);
518 eprintln!("10k point updates: {:?}", update_time);
519 eprintln!("10k prefix queries: {:?}", query_time);
520 eprintln!("10k batch updates: {:?}", batch_time);
521 eprintln!("Query p50 (approx): {:?}", query_time / 10_000);
522
523 assert!(
525 query_time < std::time::Duration::from_millis(50),
526 "10k queries too slow: {query_time:?}"
527 );
528 assert!(
529 build_time < std::time::Duration::from_millis(100),
530 "build too slow: {build_time:?}"
531 );
532 }
533
534 #[test]
539 fn single_element_tree() {
540 let mut ft = FenwickTree::new(1);
541 assert_eq!(ft.len(), 1);
542 assert!(!ft.is_empty());
543 assert_eq!(ft.total(), 0);
544 ft.update(0, 42);
545 assert_eq!(ft.get(0), 42);
546 assert_eq!(ft.prefix(0), 42);
547 assert_eq!(ft.total(), 42);
548 }
549
550 #[test]
551 fn from_values_empty() {
552 let ft = FenwickTree::from_values(&[]);
553 assert!(ft.is_empty());
554 assert_eq!(ft.len(), 0);
555 assert_eq!(ft.total(), 0);
556 }
557
558 #[test]
559 fn from_values_single() {
560 let ft = FenwickTree::from_values(&[99]);
561 assert_eq!(ft.len(), 1);
562 assert_eq!(ft.get(0), 99);
563 assert_eq!(ft.total(), 99);
564 }
565
566 #[test]
567 fn update_last_element() {
568 let mut ft = FenwickTree::new(5);
569 ft.update(4, 100);
570 assert_eq!(ft.get(4), 100);
571 assert_eq!(ft.total(), 100);
572 }
573
574 #[test]
575 fn update_negative_delta() {
576 let mut ft = FenwickTree::from_values(&[10, 20, 30]);
577 ft.update(1, -5);
578 assert_eq!(ft.get(1), 15);
579 assert_eq!(ft.total(), 55);
580 }
581
582 #[test]
583 fn update_wraps_below_zero() {
584 let mut ft = FenwickTree::from_values(&[5]);
585 ft.update(0, -10);
586 let expected = 5u32.wrapping_add((-10i32) as u32);
588 assert_eq!(ft.get(0), expected);
589 }
590
591 #[test]
592 fn get_at_zero_single_element() {
593 let ft = FenwickTree::from_values(&[42]);
594 assert_eq!(ft.get(0), 42);
595 }
596
597 #[test]
598 fn get_after_multiple_updates() {
599 let mut ft = FenwickTree::new(3);
600 ft.update(1, 10);
601 ft.update(1, 20);
602 ft.update(1, -5);
603 assert_eq!(ft.get(1), 25);
604 }
605
606 #[test]
607 fn prefix_zero_after_update() {
608 let mut ft = FenwickTree::new(5);
609 ft.update(0, 7);
610 assert_eq!(ft.prefix(0), 7);
611 }
612
613 #[test]
614 fn range_single_element() {
615 let ft = FenwickTree::from_values(&[10, 20, 30]);
616 assert_eq!(ft.range(1, 1), 20);
617 }
618
619 #[test]
620 fn range_full_equals_total() {
621 let ft = FenwickTree::from_values(&[1, 2, 3, 4, 5]);
622 assert_eq!(ft.range(0, 4), ft.total());
623 }
624
625 #[test]
626 fn range_all_zeros() {
627 let ft = FenwickTree::new(5);
628 assert_eq!(ft.range(0, 4), 0);
629 assert_eq!(ft.range(2, 3), 0);
630 }
631
632 #[test]
633 fn batch_update_empty() {
634 let mut ft = FenwickTree::from_values(&[1, 2, 3]);
635 ft.batch_update(&[]);
636 assert_eq!(ft.total(), 6);
637 }
638
639 #[test]
640 fn batch_update_same_index_multiple_times() {
641 let mut ft = FenwickTree::new(3);
642 ft.batch_update(&[(0, 10), (0, 20), (0, -5)]);
643 assert_eq!(ft.get(0), 25);
644 assert_eq!(ft.get(1), 0);
645 }
646
647 #[test]
648 fn rebuild_same_values_idempotent() {
649 let values = vec![5, 10, 15];
650 let mut ft = FenwickTree::from_values(&values);
651 let total_before = ft.total();
652 ft.rebuild(&values);
653 assert_eq!(ft.total(), total_before);
654 for (i, &v) in values.iter().enumerate() {
655 assert_eq!(ft.get(i), v);
656 }
657 }
658
659 #[test]
660 fn rebuild_all_zeros() {
661 let mut ft = FenwickTree::from_values(&[10, 20, 30]);
662 ft.rebuild(&[0, 0, 0]);
663 assert_eq!(ft.total(), 0);
664 assert_eq!(ft.get(0), 0);
665 assert_eq!(ft.get(1), 0);
666 assert_eq!(ft.get(2), 0);
667 }
668
669 #[test]
670 fn find_prefix_empty_tree() {
671 let ft = FenwickTree::new(0);
672 assert_eq!(ft.find_prefix(0), None);
673 assert_eq!(ft.find_prefix(100), None);
674 }
675
676 #[test]
677 fn find_prefix_single_element() {
678 let ft = FenwickTree::from_values(&[10]);
679 assert_eq!(ft.find_prefix(0), None);
680 assert_eq!(ft.find_prefix(10), Some(0));
681 assert_eq!(ft.find_prefix(100), Some(0));
682 }
683
684 #[test]
685 fn find_prefix_target_exceeds_total() {
686 let ft = FenwickTree::from_values(&[1, 2, 3]);
687 assert_eq!(ft.find_prefix(1000), Some(2));
689 }
690
691 #[test]
692 fn find_prefix_target_equals_total() {
693 let ft = FenwickTree::from_values(&[5, 5, 5]);
694 assert_eq!(ft.find_prefix(15), Some(2));
696 }
697
698 #[test]
699 fn find_prefix_exact_boundaries() {
700 let ft = FenwickTree::from_values(&[10, 10, 10]);
701 assert_eq!(ft.find_prefix(10), Some(0));
703 assert_eq!(ft.find_prefix(20), Some(1));
704 assert_eq!(ft.find_prefix(30), Some(2));
705 }
706
707 #[test]
708 fn resize_to_zero() {
709 let mut ft = FenwickTree::from_values(&[1, 2, 3]);
710 ft.resize(0);
711 assert!(ft.is_empty());
712 assert_eq!(ft.total(), 0);
713 }
714
715 #[test]
716 fn resize_same_size_noop() {
717 let mut ft = FenwickTree::from_values(&[1, 2, 3]);
718 ft.resize(3);
719 assert_eq!(ft.len(), 3);
720 assert_eq!(ft.total(), 6);
721 }
722
723 #[test]
724 fn resize_grow_from_zero() {
725 let mut ft = FenwickTree::new(0);
726 ft.resize(3);
727 assert_eq!(ft.len(), 3);
728 assert_eq!(ft.total(), 0);
729 ft.update(0, 5);
730 assert_eq!(ft.get(0), 5);
731 }
732
733 #[test]
734 fn clone_independence() {
735 let mut ft = FenwickTree::from_values(&[1, 2, 3]);
736 let cloned = ft.clone();
737 ft.update(0, 100);
738 assert_eq!(cloned.get(0), 1);
740 assert_eq!(ft.get(0), 101);
741 }
742
743 #[test]
744 fn debug_format_contains_name() {
745 let ft = FenwickTree::new(3);
746 let dbg = format!("{ft:?}");
747 assert!(dbg.contains("FenwickTree"));
748 }
749
750 #[test]
751 fn set_to_zero() {
752 let mut ft = FenwickTree::from_values(&[10, 20, 30]);
753 ft.set(1, 0);
754 assert_eq!(ft.get(1), 0);
755 assert_eq!(ft.total(), 40);
756 }
757
758 #[test]
759 fn set_same_value_is_noop() {
760 let mut ft = FenwickTree::from_values(&[10, 20, 30]);
761 ft.set(1, 20);
762 assert_eq!(ft.get(1), 20);
763 assert_eq!(ft.total(), 60);
764 }
765
766 #[test]
767 fn set_to_max_u32() {
768 let mut ft = FenwickTree::from_values(&[0, 0, 0]);
769 ft.set(0, u32::MAX);
770 assert_eq!(ft.get(0), u32::MAX);
771 }
772
773 #[test]
774 fn total_on_single_element() {
775 let ft = FenwickTree::from_values(&[42]);
776 assert_eq!(ft.total(), 42);
777 }
778
779 #[test]
780 fn from_values_preserves_all() {
781 let values: Vec<u32> = (1..=20).collect();
782 let ft = FenwickTree::from_values(&values);
783 for (i, &v) in values.iter().enumerate() {
784 assert_eq!(ft.get(i), v, "mismatch at index {i}");
785 }
786 assert_eq!(ft.total(), 210); }
788
789 #[test]
790 fn range_first_to_middle() {
791 let ft = FenwickTree::from_values(&[10, 20, 30, 40, 50]);
792 assert_eq!(ft.range(0, 2), 60);
793 }
794
795 #[test]
796 fn range_middle_to_end() {
797 let ft = FenwickTree::from_values(&[10, 20, 30, 40, 50]);
798 assert_eq!(ft.range(3, 4), 90);
799 }
800
801 #[test]
802 #[should_panic(expected = "out of bounds")]
803 fn update_out_of_bounds_panics() {
804 let mut ft = FenwickTree::new(3);
805 ft.update(3, 1);
806 }
807
808 #[test]
809 #[should_panic(expected = "out of bounds")]
810 fn prefix_out_of_bounds_panics() {
811 let ft = FenwickTree::new(3);
812 ft.prefix(3);
813 }
814
815 #[test]
816 #[should_panic(expected = "left")]
817 fn range_left_greater_than_right_panics() {
818 let ft = FenwickTree::from_values(&[1, 2, 3]);
819 ft.range(2, 1);
820 }
821
822 #[test]
823 #[should_panic(expected = "rebuild size mismatch")]
824 fn rebuild_wrong_size_panics() {
825 let mut ft = FenwickTree::new(3);
826 ft.rebuild(&[1, 2]);
827 }
828
829 #[test]
832 fn lowbit_correctness() {
833 assert_eq!(lowbit(1), 1);
834 assert_eq!(lowbit(2), 2);
835 assert_eq!(lowbit(3), 1);
836 assert_eq!(lowbit(4), 4);
837 assert_eq!(lowbit(6), 2);
838 assert_eq!(lowbit(8), 8);
839 assert_eq!(lowbit(12), 4);
840 }
841
842 #[test]
843 fn msb_correctness() {
844 assert_eq!(most_significant_bit(0), 0);
845 assert_eq!(most_significant_bit(1), 1);
846 assert_eq!(most_significant_bit(5), 4);
847 assert_eq!(most_significant_bit(8), 8);
848 assert_eq!(most_significant_bit(100), 64);
849 assert_eq!(most_significant_bit(1000), 512);
850 }
851}