competitive/data_structure/binary_search_tree/
split.rs

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

competitive/data_structure/binary_search_tree/
split.rs
