diff options
| -rw-r--r-- | src/fenwick.rs | 14 |
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 |