competitive/data_structure/binary_search_tree/
split.rs

1use super::{
2    BstDataAccess, BstDataMutRef, BstImmutRef, BstRoot, BstSeeker, BstSpec, EqualSide,
3    data::{self},
4    seeker::{SeekByKey, SeekBySize},
5};
6use std::{
7    borrow::Borrow,
8    ops::{Bound, RangeBounds},
9};
10
11pub struct Split<'a, Spec>
12where
13    Spec: BstSpec,
14{
15    left: Option<BstRoot<Spec>>,
16    right: Option<BstRoot<Spec>>,
17    root: &'a mut Option<BstRoot<Spec>>,
18}
19
20impl<'a, Spec> Split<'a, Spec>
21where
22    Spec: BstSpec,
23{
24    pub fn new<Seek>(
25        node: &'a mut Option<BstRoot<Spec>>,
26        seeker: Seek,
27        equal_side: EqualSide,
28    ) -> Self
29    where
30        Seek: BstSeeker<Spec = Spec>,
31    {
32        let (left, right) = Spec::split(node.take(), seeker, equal_side);
33        Self {
34            left,
35            right,
36            root: node,
37        }
38    }
39
40    pub fn left(&self) -> Option<BstImmutRef<'_, Spec>> {
41        self.left.as_ref().map(|node| node.reborrow())
42    }
43
44    pub fn right(&self) -> Option<BstImmutRef<'_, Spec>> {
45        self.right.as_ref().map(|node| node.reborrow())
46    }
47
48    pub fn left_datamut(&mut self) -> Option<BstDataMutRef<'_, Spec>> {
49        self.left.as_mut().map(|node| node.borrow_datamut())
50    }
51
52    pub fn right_datamut(&mut self) -> Option<BstDataMutRef<'_, Spec>> {
53        self.right.as_mut().map(|node| node.borrow_datamut())
54    }
55
56    pub fn manually_merge<F>(&mut self, mut f: F)
57    where
58        F: FnMut(Option<BstRoot<Spec>>, Option<BstRoot<Spec>>) -> Option<BstRoot<Spec>>,
59    {
60        self.left = f(self.left.take(), self.right.take());
61    }
62}
63
64impl<'a, Spec> Drop for Split<'a, Spec>
65where
66    Spec: BstSpec,
67{
68    fn drop(&mut self) {
69        *self.root = Spec::merge(self.left.take(), self.right.take());
70    }
71}
72
73pub struct Split3<'a, Spec>
74where
75    Spec: BstSpec,
76{
77    left: Option<BstRoot<Spec>>,
78    mid: Option<BstRoot<Spec>>,
79    right: Option<BstRoot<Spec>>,
80    root: &'a mut Option<BstRoot<Spec>>,
81}
82
83impl<'a, Spec> Split3<'a, Spec>
84where
85    Spec: BstSpec,
86{
87    pub fn new<Seek1, Seek2>(
88        node: &'a mut Option<BstRoot<Spec>>,
89        start: Bound<Seek1>,
90        end: Bound<Seek2>,
91    ) -> Self
92    where
93        Seek1: BstSeeker<Spec = Spec>,
94        Seek2: BstSeeker<Spec = Spec>,
95    {
96        let (mut rest, right) = match end {
97            Bound::Included(seeker) => Spec::split(node.take(), seeker, EqualSide::Left),
98            Bound::Excluded(seeker) => Spec::split(node.take(), seeker, EqualSide::Right),
99            Bound::Unbounded => (node.take(), None),
100        };
101        let (left, mid) = match start {
102            Bound::Included(seeker) => Spec::split(rest.take(), seeker, EqualSide::Right),
103            Bound::Excluded(seeker) => Spec::split(rest.take(), seeker, EqualSide::Left),
104            Bound::Unbounded => (None, rest),
105        };
106        Self {
107            left,
108            mid,
109            right,
110            root: node,
111        }
112    }
113
114    pub fn left(&self) -> Option<BstImmutRef<'_, Spec>> {
115        self.left.as_ref().map(|node| node.reborrow())
116    }
117
118    pub fn mid(&self) -> Option<BstImmutRef<'_, Spec>> {
119        self.mid.as_ref().map(|node| node.reborrow())
120    }
121
122    pub fn right(&self) -> Option<BstImmutRef<'_, Spec>> {
123        self.right.as_ref().map(|node| node.reborrow())
124    }
125
126    pub fn left_datamut(&mut self) -> Option<BstDataMutRef<'_, Spec>> {
127        self.left.as_mut().map(|node| node.borrow_datamut())
128    }
129
130    pub fn mid_datamut(&mut self) -> Option<BstDataMutRef<'_, Spec>> {
131        self.mid.as_mut().map(|node| node.borrow_datamut())
132    }
133
134    pub fn right_datamut(&mut self) -> Option<BstDataMutRef<'_, Spec>> {
135        self.right.as_mut().map(|node| node.borrow_datamut())
136    }
137
138    pub fn split_mid<Seek>(&mut self, seeker: Seek, equal_side: EqualSide) -> Split<'_, Spec>
139    where
140        Seek: BstSeeker<Spec = Spec>,
141    {
142        Split::new(&mut self.mid, seeker, equal_side)
143    }
144
145    pub fn manually_merge<F>(&mut self, mut f: F)
146    where
147        F: FnMut(Option<BstRoot<Spec>>, Option<BstRoot<Spec>>) -> Option<BstRoot<Spec>>,
148    {
149        let rest = f(self.mid.take(), self.right.take());
150        self.mid = f(self.left.take(), rest);
151    }
152
153    pub fn seek_by_key<K, Q, R>(node: &'a mut Option<BstRoot<Spec>>, range: R) -> Self
154    where
155        Spec: BstSpec<Data: BstDataAccess<data::marker::Key, Value = K>>,
156        K: Borrow<Q>,
157        Q: Ord + ?Sized,
158        R: RangeBounds<Q>,
159    {
160        let start = match range.start_bound() {
161            Bound::Included(key) => Bound::Included(SeekByKey::new(key)),
162            Bound::Excluded(key) => Bound::Excluded(SeekByKey::new(key)),
163            Bound::Unbounded => Bound::Unbounded,
164        };
165        let end = match range.end_bound() {
166            Bound::Included(key) => Bound::Included(SeekByKey::new(key)),
167            Bound::Excluded(key) => Bound::Excluded(SeekByKey::new(key)),
168            Bound::Unbounded => Bound::Unbounded,
169        };
170        Self::new(node, start, end)
171    }
172
173    pub fn seek_by_size<R>(node: &'a mut Option<BstRoot<Spec>>, range: R) -> Self
174    where
175        Spec: BstSpec<Data: BstDataAccess<data::marker::Size, Value = usize>>,
176        R: RangeBounds<usize>,
177    {
178        let start = match range.start_bound() {
179            Bound::Included(&index) => Bound::Included(SeekBySize::new(index)),
180            Bound::Excluded(&index) => Bound::Excluded(SeekBySize::new(index)),
181            Bound::Unbounded => Bound::Unbounded,
182        };
183        let end = match range.end_bound() {
184            Bound::Included(&index) => Bound::Included(SeekBySize::new(index)),
185            Bound::Excluded(&index) => Bound::Excluded(SeekBySize::new(index)),
186            Bound::Unbounded => Bound::Unbounded,
187        };
188        Self::new(node, start, end)
189    }
190}
191
192impl<'a, Spec> Drop for Split3<'a, Spec>
193where
194    Spec: BstSpec,
195{
196    fn drop(&mut self) {
197        let rest = Spec::merge(self.mid.take(), self.right.take());
198        *self.root = Spec::merge(self.left.take(), rest);
199    }
200}
competitive/data_structure/binary_search_tree/split.rs

competitive/data_structure/binary_search_tree/
split.rs
