cp_library_rs/data_structure/
indexedset.rs

1//! インデックス付きの集合
2
3use std::iter::FromIterator;
4use std::mem::{replace, swap};
5use std::{cmp::Ordering, fmt::Debug};
6
7/// Node
8#[derive(Debug, Clone)]
9pub struct Node<T: Ord> {
10    pub key: T,
11    pub left: Option<Box<Node<T>>>,
12    pub right: Option<Box<Node<T>>>,
13    /// 部分木のサイズ
14    pub size: usize,
15}
16
17impl<T: Ord> Node<T> {
18    pub fn new(key: T) -> Self {
19        Self {
20            key,
21            left: None,
22            right: None,
23            size: 1,
24        }
25    }
26}
27
28/// IndexedSet
29/// - スプレー木のクラス
30pub struct IndexedSet<T: Ord> {
31    size: usize,
32    pub root: Option<Box<Node<T>>>,
33}
34
35impl<T> IndexedSet<T>
36where
37    T: Ord + Clone,
38{
39    /// `a <= b`の値を返す
40    #[inline]
41    fn le(a: &T, b: &T) -> bool {
42        matches!(a.cmp(b), Ordering::Less | Ordering::Equal)
43    }
44
45    /// `a < b`の値を返す
46    #[inline]
47    fn lt(a: &T, b: &T) -> bool {
48        matches!(a.cmp(b), Ordering::Less)
49    }
50
51    /// `a >= b`の値を返す
52    #[inline]
53    fn ge(a: &T, b: &T) -> bool {
54        matches!(a.cmp(b), Ordering::Equal | Ordering::Greater)
55    }
56
57    /// `a > b`の値を返す
58    #[inline]
59    fn gt(a: &T, b: &T) -> bool {
60        matches!(a.cmp(b), Ordering::Greater)
61    }
62
63    pub fn new() -> Self {
64        Self {
65            size: 0,
66            root: None,
67        }
68    }
69
70    pub fn len(&self) -> usize {
71        self.size
72    }
73
74    pub fn is_empty(&self) -> bool {
75        self.size == 0
76    }
77
78    /// get
79    /// - 値の検索を行う
80    ///
81    /// **戻り値**
82    /// - `Option<&T>`: キーに紐づいた値
83    pub fn get(&mut self, key: &T) -> Option<&T> {
84        let lb = self.lower_bound(key);
85        if lb.is_some_and(|k| k == key) {
86            lb
87        } else {
88            None
89        }
90    }
91
92    /// 値の挿入を行う。
93    ///
94    /// **戻り値**
95    /// - `bool`: 挿入が行われたか
96    pub fn insert(&mut self, key: T) -> Option<T> {
97        // rootの取り出し
98        let root = self.root.take();
99        // splay操作（一番右の要素）
100        let (mut tmp_root, _) = splay(root, &key, Self::le);
101        if tmp_root.is_some() && tmp_root.as_ref().unwrap().key == key {
102            self.root = tmp_root;
103            let res = replace(&mut self.root.as_deref_mut().unwrap().key, key);
104            return Some(res);
105        }
106        self.root = Some(Box::new(Node::new(key.clone())));
107        if tmp_root.is_some() {
108            match key.cmp(&tmp_root.as_ref().unwrap().key) {
109                Ordering::Less | Ordering::Equal => {
110                    let mut new_left = tmp_root.as_mut().unwrap().left.take();
111                    update_size(&mut tmp_root);
112                    swap(&mut self.root.as_mut().unwrap().left, &mut new_left);
113                    swap(&mut self.root.as_mut().unwrap().right, &mut tmp_root);
114                }
115                Ordering::Greater => {
116                    let mut new_right = tmp_root.as_mut().unwrap().right.take();
117                    update_size(&mut tmp_root);
118                    swap(&mut self.root.as_mut().unwrap().right, &mut new_right);
119                    swap(&mut self.root.as_mut().unwrap().left, &mut tmp_root);
120                }
121            }
122        }
123        // 部分木のサイズを更新
124        update_size(&mut self.root);
125        // 要素数の更新
126        self.size += 1;
127        None
128    }
129
130    /// remove
131    /// - 値の削除
132    ///
133    /// **戻り値**
134    /// - `Option<T>`: 削除された値
135    pub fn remove(&mut self, key: &T) -> Option<T> {
136        if self.is_empty() {
137            return None;
138        }
139        // rootの取り出し
140        let root = self.root.take();
141        // splay操作
142        // tmp_root := keyより真に大きいノードのうち最小のもの
143        let (mut tmp_root, _) = splay(root, key, Self::le);
144        // 値が存在しないとき
145        if tmp_root.is_none() || &tmp_root.as_ref().unwrap().key != key {
146            self.root = tmp_root;
147            return None;
148        }
149        // 削除
150        if tmp_root.as_ref().unwrap().left.is_none() {
151            swap(&mut self.root, &mut tmp_root.as_mut().unwrap().right);
152        } else {
153            let root_left = tmp_root.as_mut().unwrap().left.take();
154            // 左の子のうち最大の要素を新しい根に
155            swap(&mut self.root, &mut splay(root_left, key, Self::lt).0);
156            // 根の右側に子を付け替える
157            swap(
158                &mut self.root.as_mut().unwrap().right,
159                &mut tmp_root.as_mut().unwrap().right,
160            );
161        }
162        // 部分木のサイズを更新
163        update_size(&mut self.root);
164        // 要素数の更新
165        self.size -= 1;
166        let deleted = tmp_root.take();
167        Some(deleted.unwrap().key)
168    }
169
170    /// contains_key
171    /// - 値`key`を含むか
172    pub fn contains_key(&mut self, key: &T) -> bool {
173        self.get(key).is_some_and(|k| k == key)
174    }
175
176    /// lower_bound
177    /// - `key`以上の最小の値を返す
178    pub fn lower_bound(&mut self, key: &T) -> Option<&T> {
179        // 根の取り出し
180        let root = self.root.take();
181        // スプレー操作
182        let (new_root, is_found) = splay(root, key, Self::le);
183        self.root = new_root;
184        if is_found {
185            Some(&self.root.as_ref().unwrap().key)
186        } else {
187            None
188        }
189    }
190
191    /// upper_bound
192    /// - `key`より大きい最小の値を返す
193    pub fn upper_bound(&mut self, key: &T) -> Option<&T> {
194        // 根の取り出し
195        let root = self.root.take();
196        // スプレー操作
197        let (new_root, is_found) = splay(root, key, Self::lt);
198        self.root = new_root;
199        if is_found {
200            Some(&self.root.as_ref().unwrap().key)
201        } else {
202            None
203        }
204    }
205
206    /// lower_bound_rev
207    /// - `key`以下の最大の値を返す
208    pub fn lower_bound_rev(&mut self, key: &T) -> Option<&T> {
209        // 根の取り出し
210        let root = self.root.take();
211        // スプレー操作
212        let (new_root, is_found) = splay_rev(root, key, Self::ge);
213        self.root = new_root;
214        if is_found {
215            Some(&self.root.as_ref().unwrap().key)
216        } else {
217            None
218        }
219    }
220
221    /// upper_bound_rev
222    /// - `key`未満の最大の値を返す
223    pub fn upper_bound_rev(&mut self, key: &T) -> Option<&T> {
224        // 根の取り出し
225        let root = self.root.take();
226        // スプレー操作
227        let (new_root, is_found) = splay_rev(root, key, Self::gt);
228        self.root = new_root;
229        if is_found {
230            Some(&self.root.as_ref().unwrap().key)
231        } else {
232            None
233        }
234    }
235
236    /// get_by_index
237    /// - 先頭からn番目の値を取得する（0-indexed）
238    pub fn get_by_index(&self, n: usize) -> Option<&T> {
239        if n > self.size {
240            None
241        } else {
242            get_nth(&self.root, n + 1)
243        }
244    }
245
246    /// index
247    /// - 要素`key`のインデックスを取得する（0-indexed）
248    pub fn index(&mut self, key: &T) -> Option<usize> {
249        // keyでsplayを行う
250        if self.get(key).is_some() {
251            let left_size = self
252                .root
253                .as_ref()
254                .unwrap()
255                .left
256                .as_ref()
257                .map_or(0, |node| node.size);
258            Some(left_size)
259        } else {
260            None
261        }
262    }
263}
264
265/// get_nth
266fn get_nth<T: Ord>(root: &Option<Box<Node<T>>>, n: usize) -> Option<&T> {
267    if let Some(root) = root {
268        let left_size = root.left.as_ref().map_or(0, |node| node.size);
269        match n.cmp(&(left_size + 1)) {
270            Ordering::Less => get_nth(&root.left, n),
271            Ordering::Equal => Some(&root.key),
272            Ordering::Greater => get_nth(&root.right, n - left_size - 1),
273        }
274    } else {
275        None
276    }
277}
278
279/// splay
280/// 比較関数`compare`を引数にとり、条件を満たす最小のノードを返す
281fn splay<T, C>(mut root: Option<Box<Node<T>>>, key: &T, compare: C) -> (Option<Box<Node<T>>>, bool)
282where
283    T: Ord,
284    C: Fn(&T, &T) -> bool,
285{
286    if root.is_none() {
287        return (root, false);
288    }
289    if compare(key, &root.as_ref().unwrap().key) {
290        let left = &mut root.as_mut().unwrap().left;
291        if left.is_none() {
292            return (root, true);
293        }
294        if compare(key, &left.as_ref().unwrap().key) {
295            let leftleft = left.as_mut().unwrap().left.take();
296            let (mut tmp, is_found) = splay(leftleft, key, compare);
297            // 戻す
298            swap(&mut left.as_mut().unwrap().left, &mut tmp);
299            // 親を右に回転
300            let tmp_left = rotate_right(root);
301            if !is_found {
302                return (tmp_left, true);
303            }
304            // さらに右回転
305            (rotate_right(tmp_left), true)
306        } else {
307            let leftright = left.as_mut().unwrap().right.take();
308            let (mut new_leftright, is_found) = splay(leftright, key, compare);
309            // 戻す
310            swap(&mut left.as_mut().unwrap().right, &mut new_leftright);
311            // root->left->rightがNoneでないとき
312            if !is_found {
313                return (root, true);
314            }
315            // 左の子を左回転
316            let left = root.as_mut().unwrap().left.take();
317            let mut tmp_child = rotate_left(left);
318            swap(&mut root.as_mut().unwrap().left, &mut tmp_child);
319            // 親を右回転
320            (rotate_right(root), true)
321        }
322    } else {
323        let right = &mut root.as_mut().unwrap().right;
324        if right.is_none() {
325            return (root, false);
326        }
327        if compare(key, &right.as_ref().unwrap().key) {
328            let rightleft = right.as_mut().unwrap().left.take();
329            let (mut tmp, is_found) = splay(rightleft, key, compare);
330            // 戻す
331            swap(&mut right.as_mut().unwrap().left, &mut tmp);
332            if is_found {
333                // 右の子を右回転
334                let right = root.as_mut().unwrap().right.take();
335                let mut tmp_child = rotate_right(right);
336                swap(&mut root.as_mut().unwrap().right, &mut tmp_child);
337            }
338            // 親を左回転
339            (rotate_left(root), true)
340        } else {
341            let rightright = right.as_mut().unwrap().right.take();
342            let (mut tmp, is_found) = splay(rightright, key, compare);
343            // 戻す
344            swap(&mut right.as_mut().unwrap().right, &mut tmp);
345            // 親を左回転
346            let tmp_child = rotate_left(root);
347            // さらに左回転
348            (rotate_left(tmp_child), is_found)
349        }
350    }
351}
352
353/// splay_rev
354/// - 比較関数`compare`を引数にとり、条件を満たす最小のノードを返す
355/// - splayの逆向き
356fn splay_rev<T, C>(
357    mut root: Option<Box<Node<T>>>,
358    key: &T,
359    compare: C,
360) -> (Option<Box<Node<T>>>, bool)
361where
362    T: Ord,
363    C: Fn(&T, &T) -> bool,
364{
365    if root.is_none() {
366        return (root, false);
367    }
368    if compare(key, &root.as_ref().unwrap().key) {
369        let right = &mut root.as_mut().unwrap().right;
370        if right.is_none() {
371            return (root, true);
372        }
373        if compare(key, &right.as_ref().unwrap().key) {
374            let rightright = right.as_mut().unwrap().right.take();
375            let (mut tmp, is_found) = splay_rev(rightright, key, compare);
376            // 戻す
377            swap(&mut right.as_mut().unwrap().right, &mut tmp);
378            // 親を左に回転
379            let tmp_right = rotate_left(root);
380            if !is_found {
381                return (tmp_right, true);
382            }
383            // さらに左回転
384            (rotate_left(tmp_right), true)
385        } else {
386            let rightleft = right.as_mut().unwrap().left.take();
387            let (mut new_rightleft, is_found) = splay_rev(rightleft, key, compare);
388            // 戻す
389            swap(&mut right.as_mut().unwrap().left, &mut new_rightleft);
390            // root->right->leftがNoneでないとき
391            if !is_found {
392                return (root, true);
393            }
394            // 右の子を右回転
395            let right = root.as_mut().unwrap().right.take();
396            let mut tmp_child = rotate_right(right);
397            swap(&mut root.as_mut().unwrap().right, &mut tmp_child);
398            // 親を左回転
399            (rotate_left(root), true)
400        }
401    } else {
402        let left = &mut root.as_mut().unwrap().left;
403        if left.is_none() {
404            return (root, false);
405        }
406        if compare(key, &left.as_ref().unwrap().key) {
407            let leftright = left.as_mut().unwrap().right.take();
408            let (mut tmp, is_found) = splay_rev(leftright, key, compare);
409            // 戻す
410            swap(&mut left.as_mut().unwrap().right, &mut tmp);
411            if is_found {
412                // 左の子を左回転
413                let left = root.as_mut().unwrap().left.take();
414                let mut tmp_child = rotate_left(left);
415                swap(&mut root.as_mut().unwrap().left, &mut tmp_child);
416            }
417            // 親を右回転
418            (rotate_right(root), true)
419        } else {
420            let leftleft = left.as_mut().unwrap().left.take();
421            let (mut tmp, is_found) = splay_rev(leftleft, key, compare);
422            // 戻す
423            swap(&mut left.as_mut().unwrap().left, &mut tmp);
424            // 親を右回転
425            let tmp_child = rotate_right(root);
426            // さらに右回転
427            (rotate_right(tmp_child), is_found)
428        }
429    }
430}
431
432/// 部分木のサイズを更新する
433fn update_size<T: Ord>(node: &mut Option<Box<Node<T>>>) {
434    if let Some(node) = node {
435        let left_size = node.left.as_ref().map_or(0, |node| node.size);
436        let right_size = node.right.as_ref().map_or(0, |node| node.size);
437        node.size = left_size + right_size + 1;
438    }
439}
440
441/// 右回転
442/// ```not-rust
443///        Y                      X
444///       / \       right        / \
445///      X   C  === rotate ==>  A   Y
446///     / \                        / \
447///    A   B                      B   C
448/// ```
449fn rotate_right<T: Ord>(root: Option<Box<Node<T>>>) -> Option<Box<Node<T>>> {
450    let mut root = root?;
451    let Some(mut new_root) = root.left else {
452        return Some(root);
453    };
454    root.left = new_root.right;
455    new_root.right = Some(root);
456    update_size(&mut new_root.right);
457    let mut res = Some(new_root);
458    update_size(&mut res);
459    res
460}
461
462/// 左回転
463/// ```not-rust
464///      X                          Y
465///     / \         left           / \
466///    A   Y    === rotate ==>    X   C
467///       / \                    / \
468///      B   C                  A   B
469/// ```
470fn rotate_left<T: Ord>(root: Option<Box<Node<T>>>) -> Option<Box<Node<T>>> {
471    let mut root = root?;
472    let Some(mut new_root) = root.right else {
473        return Some(root);
474    };
475    root.right = new_root.left;
476    new_root.left = Some(root);
477    update_size(&mut new_root.left);
478    let mut res = Some(new_root);
479    update_size(&mut res);
480    res
481}
482
483// ----- FromIterator -----
484impl<T: Ord + Clone> FromIterator<T> for IndexedSet<T> {
485    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
486        let mut res = IndexedSet::new();
487        for item in iter {
488            res.insert(item);
489        }
490        res
491    }
492}
493
494// ----- iterator -----
495pub struct SplayTreeIterator<'a, T: 'a + Ord> {
496    unvisited: Vec<&'a Node<T>>,
497}
498
499impl<'a, T: Ord> SplayTreeIterator<'a, T> {
500    fn push_left_edge(&mut self, mut tree: &'a Option<Box<Node<T>>>) {
501        while let Some(node) = tree.as_deref() {
502            self.unvisited.push(node);
503            tree = &node.left;
504        }
505    }
506}
507
508impl<'a, T: Ord> Iterator for SplayTreeIterator<'a, T> {
509    type Item = &'a T;
510
511    fn next(&mut self) -> Option<Self::Item> {
512        let node = self.unvisited.pop()?;
513
514        self.push_left_edge(&node.right);
515
516        Some(&node.key)
517    }
518}
519
520impl<T: Ord> IndexedSet<T> {
521    pub fn iter(&self) -> SplayTreeIterator<'_, T> {
522        let mut iter = SplayTreeIterator { unvisited: vec![] };
523        iter.push_left_edge(&self.root);
524        iter
525    }
526}
527
528impl<'a, T: Ord> IntoIterator for &'a IndexedSet<T> {
529    type IntoIter = SplayTreeIterator<'a, T>;
530    type Item = &'a T;
531
532    fn into_iter(self) -> Self::IntoIter {
533        self.iter()
534    }
535}
536
537impl<T: Ord + Debug> Debug for IndexedSet<T> {
538    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
539        f.debug_set().entries(self.iter()).finish()
540    }
541}
cp_library_rs/data_structure/indexedset.rs

cp_library_rs/data_structure/
indexedset.rs
