Rust实现线段树和懒标记

参考各家代码，用Rust实现了线段树和懒标记。

由于使用了泛型，很多操作都要用闭包自定义实现。

看代码。

// 线段树定义
pub struct SegmentTree<T: Clone>
{
    pub data: Vec<T>,
    tree: Vec<Option<T>>,
    marker: Vec<T>,                               //懒标记。
    query_op: Box<dyn Fn(T, T) -> T>, //查询时，对所有查询元素做的操作。比如加法，就是求区间的所有元素的和。
    marker_marker_op: Box<dyn Fn(T, T) -> T>, //marker加到marker上时，对marker的操作。通常我们要marker[i] += p; 来更新标记，但是泛型实现不了，并且考虑到有些用户有别的需求，所以用闭包包装。
    marker_t_op: Box<dyn Fn(T, T) -> T>, //marker应用到T时，对T的操作。考虑到有些用户有别的需求，所以用闭包包装。
    marker_mul_usize: Box<dyn Fn(T, usize) -> T>, //marker乘usize的方法。这个没法通过要求满足Mul trait自动实现。由于使用了泛型，连乘法都要交给闭包实现。。。
}

impl<T: Clone + Default + Copy + PartialEq> SegmentTree<T> {
    pub fn new(
        data: Vec<T>,
        query_op: Box<dyn Fn(T, T) -> T>,
        marker_marker_op: Box<dyn Fn(T, T) -> T>,
        marker_t_op: Box<dyn Fn(T, T) -> T>,
        marker_mul_usize: Box<dyn Fn(T, usize) -> T>,
    ) -> Self {
        let data_len = data.len();
        let mut tr = Self {
            data,
            marker: vec![T::default(); 4 * data_len], //四倍原数据大小
            tree: vec![None; 4 * data_len],           //四倍原数据大小
            query_op,
            marker_marker_op,
            marker_t_op,
            marker_mul_usize,
        };
        tr.build();
        tr
    }

    #[inline]
    pub fn get(&self, index: usize) -> Option<&T> {
        self.data.get(index)
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.data.len()
    }

    #[inline]
    fn left_child(index: usize) -> usize {
        2 * index + 1
    }

    #[inline]
    fn right_child(index: usize) -> usize {
        2 * index + 2
    }

    #[inline]
    fn build(&mut self) {
        self.build_segment_tree(0, 0, self.data.len() - 1);
    }

    // 递归Build
    fn build_segment_tree(&mut self, tree_index: usize, left: usize, right: usize) {
        if left == right {
            self.tree[tree_index] = Some(self.data[left]);
            return;
        }
        let left_tree_index = Self::left_child(tree_index);
        let right_tree_index = Self::right_child(tree_index);
        let mid = (right - left) / 2 + left;
        self.build_segment_tree(left_tree_index, left, mid);
        self.build_segment_tree(right_tree_index, mid + 1, right);
        // 左右子树数据处理方式
        if let Some(l) = self.tree[left_tree_index] {
            if let Some(r) = self.tree[right_tree_index] {
                self.tree[tree_index] = Some((self.query_op)(l, r))
            }
        }
    }

    // 返回对线段树的全部元素做query_op操作的结果
    #[inline]
    pub fn query_all(&mut self) -> T {
        self.recursion_query(0, self.data.len() - 1, 0, 0, self.data.len() - 1)
    }

    // 返回对线段树的[l..r]范围全部元素做query_op操作的结果
    pub fn query(&mut self, l: usize, r: usize) -> Result<T, &'static str> {
        if l > self.data.len() || r > self.data.len() || l > r {
            return Err("索引错误");
        }
        if l == r {
            return Ok(self.data[l]);
        }
        Ok(self.recursion_query(l, r, 0, 0, self.data.len() - 1))
    }

    // 在index表示的[current_left,current_right]范围中查询[l..r]值
    fn recursion_query(
        &mut self,
        l: usize,
        r: usize,
        index: usize,
        current_left: usize,
        current_right: usize,
    ) -> T {
        if l > current_right || r < current_left {
            return T::default();
        }
        if l == current_left && r == current_right {
            if let Some(d) = self.tree[index] {
                if l == r {
                    self.data[l] = d;
                }
                return d;
            }
            return T::default();
        }
        self.push_down(index, current_right - current_left + 1);
        let mid = current_left + (current_right - current_left) / 2;
        if l >= mid + 1 {
            return self.recursion_query(l, r, Self::right_child(index), mid + 1, current_right);
        } else if r <= mid {
            return self.recursion_query(l, r, Self::left_child(index), current_left, mid);
        }
        let l_res = self.recursion_query(l, mid, Self::left_child(index), current_left, mid);
        let r_res =
            self.recursion_query(mid + 1, r, Self::right_child(index), mid + 1, current_right);
        (self.query_op)(l_res, r_res)
    }

    // 更新index为val
    pub fn set(&mut self, index: usize, val: T) -> Result<(), &'static str> {
        if index >= self.data.len() {
            return Err("索引超过线段树长度");
        }
        // 更新数据
        self.data[index] = val;
        // 递归更新树
        self.recursion_set(0, 0, self.data.len() - 1, index, val);
        Ok(())
    }

    // 递归更新树
    fn recursion_set(&mut self, index_tree: usize, l: usize, r: usize, index: usize, val: T) {
        if l == r {
            self.tree[index_tree] = Some(val);
            return;
        }
        let mid = l + (r - l) / 2;
        let left_child = Self::left_child(index_tree);
        let right_child = Self::right_child(index_tree);
        if index >= mid + 1 {
            self.recursion_set(right_child, mid + 1, r, index, val);
        } else {
            self.recursion_set(left_child, l, mid, index, val);
        }
        // 左右子树数据求和
        if let Some(l_d) = self.tree[left_child] {
            if let Some(r_d) = self.tree[right_child] {
                self.tree[index_tree] = Some((self.query_op)(l_d, r_d));
            }
        }
    }

    // 应用所有懒标记到data数组上
    #[inline]
    pub fn apply_marker_all(&mut self) {
        self.apply_marker_lr(0, self.data.len() - 1);
    }

    // 应用懒标记到[l:r]数据范围
    #[inline]
    pub fn apply_marker_lr(&mut self, l: usize, r: usize) {
        self.apply_marker(l, r, 0, 0, self.data.len() - 1);
    }

    fn apply_marker(
        &mut self,
        l: usize,
        r: usize,
        index: usize,
        current_l: usize,
        current_r: usize,
    ) {
        if current_l > r || current_r < l || r >= self.data.len() {
            return; // 区间无交集
        } else {
            // 与目标区间有交集，但不包含于其中
            if current_l == current_r {
                if let Some(d) = self.tree[index] {
                    self.data[current_l] = d;
                }
                return;
            }
            let mid = (current_l + current_r) / 2;
            self.push_down(index, current_r - current_l + 1);
            self.apply_marker(l, r, Self::left_child(index), current_l, mid); // 递归地往下寻找
            self.apply_marker(l, r, Self::right_child(index), mid + 1, current_r);
            self.tree[index] = Some((self.query_op)(
                self.tree[Self::left_child(index)].unwrap(),
                self.tree[Self::right_child(index)].unwrap(),
            ));
            // 根据子节点更新当前节点的值
        }
    }
    #[inline]
    pub fn update_interval(&mut self, l: usize, r: usize, delta: T) {
        self.update(l, r, delta, 0, 0, self.data.len() - 1);
    }

    // 传递marker到下级
    fn push_down(&mut self, index: usize, len: usize) {
        self.marker[Self::left_child(index)] =
            (self.marker_marker_op)(self.marker[index], self.marker[Self::left_child(index)]); // 标记向下传递
        self.marker[Self::right_child(index)] =
            (self.marker_marker_op)(self.marker[index], self.marker[Self::right_child(index)]);
        if self.tree[Self::left_child(index)].is_some() {
            self.tree[Self::left_child(index)] = Some((self.marker_t_op)(
                (self.marker_mul_usize)(self.marker[index], len - (len / 2)),
                self.tree[Self::left_child(index)].unwrap(),
            ));
        }
        if self.tree[Self::right_child(index)].is_some() {
            self.tree[Self::right_child(index)] = Some((self.marker_t_op)(
                (self.marker_mul_usize)(self.marker[index], len / 2),
                self.tree[Self::right_child(index)].unwrap(),
            ));
        }
        self.marker[index] = T::default(); // 清除标记
    }

    fn update(
        &mut self,
        l: usize,
        r: usize,
        delta: T,
        index: usize,
        current_l: usize,
        current_r: usize,
    ) {
        if current_l > r || current_r < l {
            return; // 区间无交集
        } else if current_l >= l && current_r <= r {
            // 当前节点对应的区间包含在目标区间中
            if self.tree[index].is_some() {
                // 更新当前区间的值
                self.tree[index] = Some((self.query_op)(
                    self.tree[index].unwrap(),
                    (self.marker_mul_usize)(delta, current_r - current_l + 1),
                ));
            }
            // 如果不是叶子节点
            if current_r > current_l {
                // 给当前区间打上标记
                self.marker[index] = (self.marker_marker_op)(delta, self.marker[index]);
            }
        } else {
            // 与目标区间有交集，但不包含于其中
            let mid = (current_l + current_r) / 2;
            self.push_down(index, current_r - current_l + 1);
            self.update(l, r, delta, Self::left_child(index), current_l, mid); // 递归地往下寻找
            self.update(l, r, delta, Self::right_child(index), mid + 1, current_r);
            self.tree[index] = Some((self.query_op)(
                self.tree[Self::left_child(index)].unwrap(),
                self.tree[Self::right_child(index)].unwrap(),
            )); // 根据子节点更新当前节点的值
        }
    }
}

fn main() {
    let mut tr: SegmentTree<i32> = SegmentTree::new(
        vec![1, 3, 4, 0, 0, 4, 5, 0],
        Box::new(|a, b| a + b),
        Box::new(|a, b| a + b),
        Box::new(|a, b| a + b),
        Box::new(|a, b| a * (b as i32)),
    );
    let _ = tr.set(1, 2); //点更新，即把data[1]设为2
    tr.update_interval(0, 2, -1); //区间更新，即[0:2]每个元素减1
    tr.update_interval(1, 3, 2); //区间更新，即[1:3]每个元素加2
    tr.apply_marker_all(); //应用全部marker到data数组
    println!("{}", tr.query_all()); //输出19，即全部元素的和
    println!("{:?}", tr.data); //输出[0, 3, 5, 2, 0, 4, 5, 0]
}

做一道题验证一下这个线段树的正确性，直接看我写的1589. 所有排列中的最大和题解即可（虽然这道题用差分数组最快，但是作为线段树验证还是很方便的）。