Optimize range_query to match fenwick.cHEADmaster

1 files changed, 12 insertions, 2 deletions

diff --git a/src/fenwick.rs b/src/fenwick.rs
index e26feff..29b435d 100644
--- a/src/fenwick.rs
+++ b/src/fenwick.rs
@@ -32,8 +32,18 @@ pub fn query(f: &[i32], mut i: usize) -> i32 {
 }
 
 // Query for sum between indices i and j inclusive
-pub fn range_query(f: &[i32], i: usize, j: usize) -> i32 {
-    query(f, j) - query(f, i - 1)
+pub fn range_query(f: &[i32], mut i: usize, mut j: usize) -> i32 {
+    let mut ret = 0;
+    while j >= i {
+        ret += f[j];
+        j -= j & j.wrapping_neg();
+    }
+    i -= 1;
+    while i > j {
+        ret -= f[i];
+        i -= i & i.wrapping_neg();
+    }
+    ret
 }
 
 // Search for largest index with prefix sum less than or equal to s