atcoder/lib/data_structure/
segment_tree.rs

1//! 点更新・区間集約を行うセグメント木です。
2
3use std::fmt::Debug;
4use std::ops::Index;
5
6/// モノイド `(T, op, identity)` 上のセグメント木です。
7///
8/// 添字は 0-indexed、区間は半開区間 `[left, right)` です。`op` は結合則を満たし、
9/// `identity` はその単位元である必要があります。`op` が非可換でも、区間内の左から右の
10/// 順序で集約します。構築は `O(N)`、点更新・区間取得・
11/// 二分探索は `O(log N)` です。
12///
13/// # Examples
14///
15/// ```
16/// use atcoder::data_structure::segment_tree::SegmentTree;
17///
18/// let mut tree = SegmentTree::from_vec([1, 2, 3, 4], |a, b| a + b, 0);
19/// assert_eq!(tree.get(1, 4), 9);
20/// tree.set(2, 10);
21/// assert_eq!(tree.get(0, 3), 13);
22/// ```
23#[derive(Clone)]
24pub struct SegmentTree<T, F> {
25    op: F,
26    identity: T,
27    len: usize,
28    size: usize,
29    data: Vec<T>,
30}
31
32impl<T: Copy, F: Fn(T, T) -> T> SegmentTree<T, F> {
33    /// 長さ `len`、各要素が `identity` の木を作ります。計算量・メモリ量は `O(N)` です。
34    ///
35    /// `op` は結合則を満たし、`identity` は左右両方の単位元である必要があります。
36    pub fn new(op: F, len: usize, identity: T) -> Self {
37        let size = len.next_power_of_two().max(1);
38        Self {
39            op,
40            identity,
41            len,
42            size,
43            data: vec![identity; size * 2],
44        }
45    }
46
47    /// `values` の配列から木を構築します。計算量・メモリ量は `O(N)` です。
48    ///
49    /// `op` は結合則を満たし、`identity` は左右両方の単位元である必要があります。
50    pub fn from_vec(values: impl Into<Vec<T>>, op: F, identity: T) -> Self {
51        let values = values.into();
52        let len = values.len();
53        let size = len.next_power_of_two().max(1);
54        let mut data = vec![identity; size * 2];
55        data[size..size + len].copy_from_slice(&values);
56        for index in (1..size).rev() {
57            data[index] = op(data[index << 1], data[index << 1 | 1]);
58        }
59        Self {
60            op,
61            identity,
62            len,
63            size,
64            data,
65        }
66    }
67
68    /// [`from_vec`](Self::from_vec) と同じです。`values` の配列から木を構築します。
69    pub fn from(op: F, values: Vec<T>, identity: T) -> Self {
70        Self::from_vec(values, op, identity)
71    }
72
73    /// 要素数を返します。
74    pub const fn len(&self) -> usize {
75        self.len
76    }
77
78    /// 要素がなければ `true` を返します。
79    pub const fn is_empty(&self) -> bool {
80        self.len == 0
81    }
82
83    /// 0 始まりの `index` 番目を `value` に更新します。計算量は `O(log N)` です。
84    ///
85    /// `index >= len` の場合は panic します。
86    pub fn set(&mut self, index: usize, value: T) {
87        assert!(index < self.len, "index out of bounds");
88        let mut index = index + self.size;
89        self.data[index] = value;
90        while index > 1 {
91            index >>= 1;
92            self.data[index] = (self.op)(self.data[index << 1], self.data[index << 1 | 1]);
93        }
94    }
95
96    /// 0 始まりの `index` 番目の現在値を `f` に渡し、戻り値で更新します。
97    ///
98    /// `f` は1回だけ呼び出し、更新後は祖先の集約値も再計算します。
99    /// 計算量は `f` の実行時間を除いて `O(log N)` です。`index >= len` なら
100    /// `f` を呼び出す前に panic します。
101    ///
102    /// ```
103    /// use atcoder::data_structure::segment_tree::SegmentTree;
104    ///
105    /// let mut seg = SegmentTree::from_vec([1, 2, 3], |a, b| a + b, 0);
106    /// seg.update(1, |v| v + 5);
107    /// seg.update(2, |v| v.max(10));
108    /// assert_eq!(seg.at(1), 7);
109    /// assert_eq!(seg.get(0, 3), 18);
110    /// ```
111    pub fn update(&mut self, index: usize, f: impl FnOnce(T) -> T) {
112        self.set(index, f(self.at(index)));
113    }
114
115    /// 0 始まりの `index` 番目の値を返します。`index >= len` なら panic します。
116    /// 計算量は `O(1)` です。
117    pub fn at(&self, index: usize) -> T {
118        assert!(index < self.len, "index out of bounds");
119        self.data[self.size + index]
120    }
121
122    /// 半開区間 `[left, right)` の集約値を返します。
123    ///
124    /// 添字は 0 始まりです。`left > right` または `right > len` の場合は panic します。
125    /// 空区間には単位元を返します。計算量は `O(log N)` です。
126    pub fn get(&self, left: usize, right: usize) -> T {
127        assert!(left <= right && right <= self.len, "range out of bounds");
128        let mut left = left + self.size;
129        let mut right = right + self.size;
130        let mut left_value = self.identity;
131        let mut right_value = self.identity;
132        while left < right {
133            if left & 1 == 1 {
134                left_value = (self.op)(left_value, self.data[left]);
135                left += 1;
136            }
137            if right & 1 == 1 {
138                right -= 1;
139                right_value = (self.op)(self.data[right], right_value);
140            }
141            left >>= 1;
142            right >>= 1;
143        }
144        (self.op)(left_value, right_value)
145    }
146
147    /// `predicate(get(left, right))` が真である最大の `right` を返します。
148    ///
149    /// `predicate(identity)` は真で、右端を伸ばしたとき述語が真から偽へ高々一度だけ
150    /// 変化する必要があります。集約は `op` の左から右の順で行います。
151    /// 返り値は `left..right` が条件を満たす最大の `right` です。`left > len` なら panic します。
152    pub fn max_right<P>(&self, mut left: usize, predicate: P) -> usize
153    where
154        P: Fn(T) -> bool,
155    {
156        assert!(left <= self.len, "index out of bounds");
157        assert!(predicate(self.identity), "predicate(identity) must be true");
158        if left == self.len {
159            return self.len;
160        }
161        left += self.size;
162        let mut value = self.identity;
163        loop {
164            while left & 1 == 0 {
165                left >>= 1;
166            }
167            let next = (self.op)(value, self.data[left]);
168            if !predicate(next) {
169                while left < self.size {
170                    left <<= 1;
171                    let next = (self.op)(value, self.data[left]);
172                    if predicate(next) {
173                        value = next;
174                        left += 1;
175                    }
176                }
177                return (left - self.size).min(self.len);
178            }
179            value = next;
180            left += 1;
181            if left.is_power_of_two() {
182                break;
183            }
184        }
185        self.len
186    }
187
188    /// `predicate(get(left, right))` が真である最小の `left` を返します。
189    ///
190    /// `predicate(identity)` は真で、左端を左へ動かしたとき述語が真から偽へ高々一度だけ
191    /// 変化する必要があります。集約は `op` の左から右の順で行います。
192    /// 返り値は `left..right` が条件を満たす最小の `left` です。`right > len` なら panic します。
193    pub fn min_left<P>(&self, mut right: usize, predicate: P) -> usize
194    where
195        P: Fn(T) -> bool,
196    {
197        assert!(right <= self.len, "index out of bounds");
198        assert!(predicate(self.identity), "predicate(identity) must be true");
199        if right == 0 {
200            return 0;
201        }
202        right += self.size;
203        let mut value = self.identity;
204        loop {
205            right -= 1;
206            while right > 1 && right % 2 == 1 {
207                right >>= 1;
208            }
209            let next = (self.op)(self.data[right], value);
210            if !predicate(next) {
211                while right < self.size {
212                    right = right << 1 | 1;
213                    let next = (self.op)(self.data[right], value);
214                    if predicate(next) {
215                        value = next;
216                        right -= 1;
217                    }
218                }
219                return (right + 1 - self.size).min(self.len);
220            }
221            value = next;
222            if right.is_power_of_two() {
223                break;
224            }
225        }
226        0
227    }
228}
229
230impl<T: Debug, F> SegmentTree<T, F> {
231    /// 葉のうち、元配列に対応する要素だけをデバッグ出力します。
232    /// 出力先や書式は [`crate::debug!`] マクロの設定に従います。
233    pub fn debug_values(&self) {
234        crate::debug!("{:?}", &self.data[self.size..self.size + self.len]);
235    }
236}
237
238impl<T, F> Index<usize> for SegmentTree<T, F> {
239    type Output = T;
240
241    fn index(&self, index: usize) -> &Self::Output {
242        assert!(index < self.len, "index out of bounds");
243        &self.data[self.size + index]
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::SegmentTree;
250
251    #[test]
252    fn gets_and_searches() {
253        let tree = SegmentTree::from_vec(vec![1, 2, 3, 4], |a, b| a + b, 0);
254        assert_eq!(tree.get(1, 3), 5);
255        assert_eq!(tree.max_right(0, |sum| sum <= 5), 2);
256        assert_eq!(tree.min_left(4, |sum| sum <= 7), 2);
257    }
258
259    #[test]
260    fn updates_values_and_ancestor_aggregates() {
261        let mut tree = SegmentTree::from_vec([1, 2, 3], |a, b| a + b, 0);
262        let mut calls = 0;
263        tree.update(0, |value| {
264            calls += 1;
265            value + 4
266        });
267        assert_eq!(calls, 1);
268        assert_eq!(tree.get(0, 2), 7);
269
270        // キャプチャを消費する FnOnce クロージャも受け取れる。
271        let candidates = vec![10];
272        tree.update(2, move |value| {
273            value.max(candidates.into_iter().next().unwrap())
274        });
275        assert_eq!(tree.at(2), 10);
276        assert_eq!(tree.get(1, 3), 12);
277        assert_eq!(tree.get(0, 3), 17);
278    }
279
280    #[test]
281    #[should_panic(expected = "index out of bounds")]
282    fn update_rejects_out_of_bounds_before_calling_closure() {
283        let mut tree = SegmentTree::from_vec([1], |a, b| a + b, 0);
284        tree.update(1, |_| panic!("closure must not be called"));
285    }
286
287    #[test]
288    fn debug_values_accepts_empty_and_non_power_of_two_trees() {
289        SegmentTree::from_vec(Vec::<i32>::new(), |a, b| a + b, 0).debug_values();
290        SegmentTree::from_vec(vec![1, 2, 3], |a, b| a + b, 0).debug_values();
291    }
292}
atcoder/lib/data_structure/segment_tree.rs

atcoder/lib/data_structure/
segment_tree.rs
