Rust实现线段树和懒标记

参考各家代码，用Rust实现了线段树和懒标记。
由于使用了泛型，很多操作都要用闭包自定义实现。
看代码。

1. // 线段树定义
2. pub struct SegmentTree<T: Clone>
3. {
4.     pub data: Vec<T>,
5.     tree: Vec<Option<T>>,
6.     marker: Vec<T>,                               //懒标记。
7.     query_op: Box<dyn Fn(T, T) -> T>, //查询时，对所有查询元素做的操作。比如加法，就是求区间的所有元素的和。
8.     marker_marker_op: Box<dyn Fn(T, T) -> T>, //marker加到marker上时，对marker的操作。通常我们要marker[i] += p; 来更新标记，但是泛型实现不了，并且考虑到有些用户有别的需求，所以用闭包包装。
9.     marker_t_op: Box<dyn Fn(T, T) -> T>, //marker应用到T时，对T的操作。考虑到有些用户有别的需求，所以用闭包包装。
10.     marker_mul_usize: Box<dyn Fn(T, usize) -> T>, //marker乘usize的方法。这个没法通过要求满足Mul trait自动实现。由于使用了泛型，连乘法都要交给闭包实现。。。
11. }
13. impl<T: Clone + Default + Copy + PartialEq> SegmentTree<T> {
14.     pub fn new(
15.         data: Vec<T>,
16.         query_op: Box<dyn Fn(T, T) -> T>,
17.         marker_marker_op: Box<dyn Fn(T, T) -> T>,
18.         marker_t_op: Box<dyn Fn(T, T) -> T>,
19.         marker_mul_usize: Box<dyn Fn(T, usize) -> T>,
20.     ) -> Self {
21.         let data_len = data.len();
22.         let mut tr = Self {
23.             data,
24.             marker: vec![T::default(); 4 * data_len], //四倍原数据大小
25.             tree: vec![None; 4 * data_len],           //四倍原数据大小
26.             query_op,
27.             marker_marker_op,
28.             marker_t_op,
29.             marker_mul_usize,
30.         };
31.         tr.build();
32.         tr
33.     }
35.     #[inline]
36.     pub fn get(&self, index: usize) -> Option<&T> {
37.         self.data.get(index)
38.     }
40.     #[inline]
41.     pub fn len(&self) -> usize {
42.         self.data.len()
43.     }
45.     #[inline]
46.     fn left_child(index: usize) -> usize {
47.         2 * index + 1
48.     }
50.     #[inline]
51.     fn right_child(index: usize) -> usize {
52.         2 * index + 2
53.     }
55.     #[inline]
56.     fn build(&mut self) {
57.         self.build_segment_tree(0, 0, self.data.len() - 1);
58.     }
60.     // 递归Build
61.     fn build_segment_tree(&mut self, tree_index: usize, left: usize, right: usize) {
62.         if left == right {
63.             self.tree[tree_index] = Some(self.data[left]);
64.             return;
65.         }
66.         let left_tree_index = Self::left_child(tree_index);
67.         let right_tree_index = Self::right_child(tree_index);
68.         let mid = (right - left) / 2 + left;
69.         self.build_segment_tree(left_tree_index, left, mid);
70.         self.build_segment_tree(right_tree_index, mid + 1, right);
71.         // 左右子树数据处理方式
72.         if let Some(l) = self.tree[left_tree_index] {
73.             if let Some(r) = self.tree[right_tree_index] {
74.                 self.tree[tree_index] = Some((self.query_op)(l, r))
75.             }
76.         }
77.     }
79.     // 返回对线段树的全部元素做query_op操作的结果
80.     #[inline]
81.     pub fn query_all(&mut self) -> T {
82.         self.recursion_query(0, self.data.len() - 1, 0, 0, self.data.len() - 1)
83.     }
85.     // 返回对线段树的[l..r]范围全部元素做query_op操作的结果
86.     pub fn query(&mut self, l: usize, r: usize) -> Result<T, &'static str> {
87.         if l > self.data.len() || r > self.data.len() || l > r {
88.             return Err("索引错误");
89.         }
90.         if l == r {
91.             return Ok(self.data[l]);
92.         }
93.         Ok(self.recursion_query(l, r, 0, 0, self.data.len() - 1))
94.     }
96.     // 在index表示的[current_left,current_right]范围中查询[l..r]值
97.     fn recursion_query(
98.         &mut self,
99.         l: usize,
100.         r: usize,
101.         index: usize,
102.         current_left: usize,
103.         current_right: usize,
104.     ) -> T {
105.         if l > current_right || r < current_left {
106.             return T::default();
107.         }
108.         if l == current_left && r == current_right {
109.             if let Some(d) = self.tree[index] {
110.                 if l == r {
111.                     self.data[l] = d;
112.                 }
113.                 return d;
114.             }
115.             return T::default();
116.         }
117.         self.push_down(index, current_right - current_left + 1);
118.         let mid = current_left + (current_right - current_left) / 2;
119.         if l >= mid + 1 {
120.             return self.recursion_query(l, r, Self::right_child(index), mid + 1, current_right);
121.         } else if r <= mid {
122.             return self.recursion_query(l, r, Self::left_child(index), current_left, mid);
123.         }
124.         let l_res = self.recursion_query(l, mid, Self::left_child(index), current_left, mid);
125.         let r_res =
126.             self.recursion_query(mid + 1, r, Self::right_child(index), mid + 1, current_right);
127.         (self.query_op)(l_res, r_res)
128.     }
130.     // 更新index为val
131.     pub fn set(&mut self, index: usize, val: T) -> Result<(), &'static str> {
132.         if index >= self.data.len() {
133.             return Err("索引超过线段树长度");
134.         }
135.         // 更新数据
136.         self.data[index] = val;
137.         // 递归更新树
138.         self.recursion_set(0, 0, self.data.len() - 1, index, val);
139.         Ok(())
140.     }
142.     // 递归更新树
143.     fn recursion_set(&mut self, index_tree: usize, l: usize, r: usize, index: usize, val: T) {
144.         if l == r {
145.             self.tree[index_tree] = Some(val);
146.             return;
147.         }
148.         let mid = l + (r - l) / 2;
149.         let left_child = Self::left_child(index_tree);
150.         let right_child = Self::right_child(index_tree);
151.         if index >= mid + 1 {
152.             self.recursion_set(right_child, mid + 1, r, index, val);
153.         } else {
154.             self.recursion_set(left_child, l, mid, index, val);
155.         }
156.         // 左右子树数据求和
157.         if let Some(l_d) = self.tree[left_child] {
158.             if let Some(r_d) = self.tree[right_child] {
159.                 self.tree[index_tree] = Some((self.query_op)(l_d, r_d));
160.             }
161.         }
162.     }
164.     // 应用所有懒标记到data数组上
165.     #[inline]
166.     pub fn apply_marker_all(&mut self) {
167.         self.apply_marker_lr(0, self.data.len() - 1);
168.     }
170.     // 应用懒标记到[l:r]数据范围
171.     #[inline]
172.     pub fn apply_marker_lr(&mut self, l: usize, r: usize) {
173.         self.apply_marker(l, r, 0, 0, self.data.len() - 1);
174.     }
176.     fn apply_marker(
177.         &mut self,
178.         l: usize,
179.         r: usize,
180.         index: usize,
181.         current_l: usize,
182.         current_r: usize,
183.     ) {
184.         if current_l > r || current_r < l || r >= self.data.len() {
185.             return; // 区间无交集
186.         } else {
187.             // 与目标区间有交集，但不包含于其中
188.             if current_l == current_r {
189.                 if let Some(d) = self.tree[index] {
190.                     self.data[current_l] = d;
191.                 }
192.                 return;
193.             }
194.             let mid = (current_l + current_r) / 2;
195.             self.push_down(index, current_r - current_l + 1);
196.             self.apply_marker(l, r, Self::left_child(index), current_l, mid); // 递归地往下寻找
197.             self.apply_marker(l, r, Self::right_child(index), mid + 1, current_r);
198.             self.tree[index] = Some((self.query_op)(
199.                 self.tree[Self::left_child(index)].unwrap(),
200.                 self.tree[Self::right_child(index)].unwrap(),
201.             ));
202.             // 根据子节点更新当前节点的值
203.         }
204.     }
205.     #[inline]
206.     pub fn update_interval(&mut self, l: usize, r: usize, delta: T) {
207.         self.update(l, r, delta, 0, 0, self.data.len() - 1);
208.     }
210.     // 传递marker到下级
211.     fn push_down(&mut self, index: usize, len: usize) {
212.         self.marker[Self::left_child(index)] =
213.             (self.marker_marker_op)(self.marker[index], self.marker[Self::left_child(index)]); // 标记向下传递
214.         self.marker[Self::right_child(index)] =
215.             (self.marker_marker_op)(self.marker[index], self.marker[Self::right_child(index)]);
216.         if self.tree[Self::left_child(index)].is_some() {
217.             self.tree[Self::left_child(index)] = Some((self.marker_t_op)(
218.                 (self.marker_mul_usize)(self.marker[index], len - (len / 2)),
219.                 self.tree[Self::left_child(index)].unwrap(),
220.             ));
221.         }
222.         if self.tree[Self::right_child(index)].is_some() {
223.             self.tree[Self::right_child(index)] = Some((self.marker_t_op)(
224.                 (self.marker_mul_usize)(self.marker[index], len / 2),
225.                 self.tree[Self::right_child(index)].unwrap(),
226.             ));
227.         }
228.         self.marker[index] = T::default(); // 清除标记
229.     }
231.     fn update(
232.         &mut self,
233.         l: usize,
234.         r: usize,
235.         delta: T,
236.         index: usize,
237.         current_l: usize,
238.         current_r: usize,
239.     ) {
240.         if current_l > r || current_r < l {
241.             return; // 区间无交集
242.         } else if current_l >= l && current_r <= r {
243.             // 当前节点对应的区间包含在目标区间中
244.             if self.tree[index].is_some() {
245.                 // 更新当前区间的值
246.                 self.tree[index] = Some((self.query_op)(
247.                     self.tree[index].unwrap(),
248.                     (self.marker_mul_usize)(delta, current_r - current_l + 1),
249.                 ));
250.             }
251.             // 如果不是叶子节点
252.             if current_r > current_l {
253.                 // 给当前区间打上标记
254.                 self.marker[index] = (self.marker_marker_op)(delta, self.marker[index]);
255.             }
256.         } else {
257.             // 与目标区间有交集，但不包含于其中
258.             let mid = (current_l + current_r) / 2;
259.             self.push_down(index, current_r - current_l + 1);
260.             self.update(l, r, delta, Self::left_child(index), current_l, mid); // 递归地往下寻找
261.             self.update(l, r, delta, Self::right_child(index), mid + 1, current_r);
262.             self.tree[index] = Some((self.query_op)(
263.                 self.tree[Self::left_child(index)].unwrap(),
264.                 self.tree[Self::right_child(index)].unwrap(),
265.             )); // 根据子节点更新当前节点的值
266.         }
267.     }
268. }
270. fn main() {
271.     let mut tr: SegmentTree<i32> = SegmentTree::new(
272.         vec![1, 3, 4, 0, 0, 4, 5, 0],
273.         Box::new(|a, b| a + b),
274.         Box::new(|a, b| a + b),
275.         Box::new(|a, b| a + b),
276.         Box::new(|a, b| a * (b as i32)),
277.     );
278.     let _ = tr.set(1, 2); //点更新，即把data[1]设为2
279.     tr.update_interval(0, 2, -1); //区间更新，即[0:2]每个元素减1
280.     tr.update_interval(1, 3, 2); //区间更新，即[1:3]每个元素加2
281.     tr.apply_marker_all(); //应用全部marker到data数组
282.     println!("{}", tr.query_all()); //输出19，即全部元素的和
283.     println!("{:?}", tr.data); //输出[0, 3, 5, 2, 0, 4, 5, 0]
284. }

复制代码

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

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！