Given an array nums
and two types of queries where you should update the value of an index in the array, and retrieve the sum of a range in the array.
Implement the NumArray
class:
NumArray(int[] nums)
Initializes the object with the integer array nums
.void update(int index, int val)
Updates the value of nums[index]
to be val
.int sumRange(int left, int right)
Returns the sum of the subarray nums[left, right]
(i.e., nums[left] + nums[left + 1], ..., nums[right]
).
Example 1:
Input ["NumArray", "sumRange", "update", "sumRange"] [[[1, 3, 5]], [0, 2], [1, 2], [0, 2]] Output [null, 9, null, 8] Explanation NumArray numArray = new NumArray([1, 3, 5]); numArray.sumRange(0, 2); // return 9 = sum([1,3,5]) numArray.update(1, 2); // nums = [1,2,5] numArray.sumRange(0, 2); // return 8 = sum([1,2,5])
Constraints:
1 <= nums.length <= 3 * 104
-100 <= nums[i] <= 100
0 <= index < nums.length
-100 <= val <= 100
0 <= left <= right < nums.length
3 * 104
calls will be made to update
and sumRange
.struct BitTree {
tree: Vec<i32>,
data: Vec<i32>,
n: usize,
}
impl BitTree {
fn new(n: usize) -> Self {
let tree = vec![0; n];
let data = vec![0; n];
let n = n;
BitTree { tree, data, n }
}
fn get(&self, i: usize) -> i32 {
self.data[i]
}
fn sum(&self, i: usize) -> i32 {
let mut res = 0;
let down_iter = std::iter::successors(Some(i), |&i| {
let j = i & (i + 1);
if j > 0 {
Some(j - 1)
} else {
None
}
});
for j in down_iter {
res += self.tree[j];
}
res
}
fn add(&mut self, i: usize, v: i32) {
self.data[i] += v;
let n = self.n;
let up_iter = std::iter::successors(Some(i), |&i| {
let j = i | (i + 1);
if j < n {
Some(j)
} else {
None
}
});
for j in up_iter {
self.tree[j] += v;
}
}
}
struct NumArray {
bit_tree: BitTree,
}
impl NumArray {
fn new(nums: Vec<i32>) -> Self {
let n = nums.len();
let mut bit_tree = BitTree::new(n);
for i in 0..n {
bit_tree.add(i, nums[i]);
}
NumArray { bit_tree }
}
fn update(&mut self, i: i32, val: i32) {
let i = i as usize;
self.bit_tree.add(i as usize, val - self.bit_tree.get(i))
}
fn sum_range(&self, i: i32, j: i32) -> i32 {
let i = i as usize;
let j = j as usize;
if i > 0 {
self.bit_tree.sum(j) - self.bit_tree.sum(i - 1)
} else {
self.bit_tree.sum(j)
}
}
}
#[test]
fn test() {
let nums = vec![1, 3, 5];
let mut obj = NumArray::new(nums);
assert_eq!(obj.sum_range(0, 2), 9);
obj.update(1, 2);
assert_eq!(obj.sum_range(0, 2), 8);
}