RUST GYM Leetcode 307 Range Sum Query

307. Range Sum Query - Mutable

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 * 10⁴
-100 <= nums[i] <= 100
0 <= index < nums.length
-100 <= val <= 100
0 <= left <= right < nums.length
At most 3 * 10⁴ calls will be made to update and sumRange.

Rust Solution

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);
}

Having problems with this solution? Click here to submit an issue on github.