307. Range Sum Query - Mutable

Given an integer array nums, handle multiple queries of the following types:

Update the value of an element in nums.
Calculate the sum of the elements of nums between indices left and right inclusive where left <= right.
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 elements of nums between indices left and right inclusive (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 1 + 3 + 5 = 9
numArray.update(1, 2); // nums = [1, 2, 5]
numArray.sumRange(0, 2); // return 1 + 2 + 5 = 8

Constraints:

• 1 <= nums.length <= 3 * 104
• -100 <= nums[i] <= 100
• 0 <= index < nums.length
• -100 <= val <= 100
• 0 <= left <= right < nums.length
• At most 3 * 104 calls will be made to update and sumRange.

---

从这一题学到了一个新的数据结构， segment tree， 但是拿数组实现的时候没有成功， 或者说是没有好的方法来构建， 于是就用最简单的二叉树形式来实现了。

segment tree 简单来说就是把一个数组不停的二分， 每个分出来的 slice 作为一个节点， 叶子节点为数组中的单个元素。每个节点可以保存加和、平均值之类的与其对应的 slice 有关的数据。 就这个题来讲， 我们在每个节点上保存了 slice 的加和， 查询的时候，如果查询的范围正好是当前节点 slice 的范围， 则我们直接返回这个加和， 否则的话继续向下查询。更新叶子节点的时候我们选择更新差值（新值-旧值）， 从根节点开始遍历，所走过的节点的 slice 只要包含该叶子节点那该节点的加和就需要更新成 sum + diff。

---

use std::cell::RefCell;
use std::rc::Rc;

#[derive(Debug)]
struct NumArray {
    nums: Vec<i32>,
    tree: Option<Rc<RefCell<Node>>>,
}

#[derive(Debug)]
struct Node {
    start: usize,
    end: usize,
    sum: i32,
    left: Option<Rc<RefCell<Node>>>,
    right: Option<Rc<RefCell<Node>>>,
}

impl Node {
    fn new(start: usize, end: usize, sum: i32) -> Self {
        Self {
            start: start,
            end: end,
            sum: sum,
            left: None,
            right: None,
        }
    }
}

fn update(root: &Option<Rc<RefCell<Node>>>, index: usize, diff: i32) {
    if let Some(node) = root {
        let mut b = node.borrow_mut();
        if b.start <= index && b.end >= index {
            b.sum += diff;
            update(&b.left, index, diff);
            update(&b.right, index, diff);
        }
    }
}

fn query(root: &Option<Rc<RefCell<Node>>>, start: usize, end: usize) -> i32 {
    if let Some(node) = root {
        let b = node.borrow();
        if start > b.end || end < b.start {
            return 0;
        }
        if b.start >= start && b.end <= end {
            return b.sum;
        }
        let left = query(&b.left, start, end);
        let right = query(&b.right, start, end);
        return left + right;
    }
    0
}

fn build_tree(nums: &Vec<i32>, start: usize, end: usize) -> Option<Rc<RefCell<Node>>> {
    if start == end {
        return Some(Rc::new(RefCell::new(Node::new(start, end, nums[start]))));
    }
    let sum: i32 = nums[start..=end].iter().map(|v| *v).sum();
    let mut node = Node::new(start, end, sum);
    let left = build_tree(nums, start, start + (end - start) / 2);
    let right = build_tree(nums, start + (end - start) / 2 + 1, end);
    node.left = left;
    node.right = right;
    Some(Rc::new(RefCell::new(node)))
}

/**
 * `&self` means the method takes an immutable reference.
 * If you need a mutable reference, change it to `&mut self` instead.
 */
impl NumArray {
    fn new(nums: Vec<i32>) -> Self {
        let tree = build_tree(&nums, 0, nums.len() - 1);
        Self { nums, tree }
    }

    fn update(&mut self, index: i32, val: i32) {
        let diff = val - self.nums[index as usize];
        update(&self.tree, index as usize, diff);
        self.nums[index as usize] = val;
    }

    fn sum_range(&self, left: i32, right: i32) -> i32 {
        query(&self.tree, left as usize, right as usize)
    }
}