general_sam/utils/
treap.rs1use std::borrow::Cow;
4use std::ops::Deref;
5use std::sync::Arc;
6
7use rand::random;
8
9pub type NeedSwap = bool;
10
11#[derive(Clone, Debug)]
12pub enum SplitTo {
13 Left,
14 Right,
15}
16
17pub trait TreapNodeData: Clone {
18 type TagType: Default;
19
20 fn get_tag(&self) -> Option<Self::TagType>;
21 fn reset_tag(&mut self);
22 fn add_tag(&mut self, tag: Self::TagType) -> NeedSwap;
23 fn update(&mut self, left: Option<&Self>, right: Option<&Self>);
24}
25
26#[derive(Clone, Debug)]
27pub struct TreapTree<DataType: TreapNodeData>(Option<Arc<TreapNode<DataType>>>);
28
29#[derive(Clone, Debug)]
30pub struct TreapNode<DataType: TreapNodeData> {
31 pub data: DataType,
32 height: u64,
33 _left: TreapTree<DataType>,
34 _right: TreapTree<DataType>,
35}
36
37impl<DataType: TreapNodeData> TreapNode<DataType> {
38 fn new(data: DataType) -> Self {
39 Self {
40 data,
41 height: random(),
42 _left: Default::default(),
43 _right: Default::default(),
44 }
45 }
46
47 fn new_from_rng<R: FnMut() -> u64>(data: DataType, mut rng: R) -> Self {
48 Self {
49 data,
50 height: rng(),
51 _left: Default::default(),
52 _right: Default::default(),
53 }
54 }
55
56 fn update(&mut self) {
57 self.data.update(
58 self._left.as_ref().map(|x| &x.data),
59 self._right.as_ref().map(|x| &x.data),
60 )
61 }
62
63 fn add_tag(&mut self, tag: DataType::TagType) {
64 if self.data.add_tag(tag) {
65 std::mem::swap(&mut self._left, &mut self._right)
66 }
67 }
68
69 #[must_use]
70 pub fn get_left(&self) -> Cow<'_, TreapTree<DataType>> {
71 match self.data.get_tag() {
72 Some(tag) => Cow::Owned(self._left.add_tag(tag)),
73 None => Cow::Borrowed(&self._left),
74 }
75 }
76
77 #[must_use]
78 pub fn get_right(&self) -> Cow<'_, TreapTree<DataType>> {
79 match self.data.get_tag() {
80 Some(tag) => Cow::Owned(self._right.add_tag(tag)),
81 None => Cow::Borrowed(&self._right),
82 }
83 }
84
85 fn set_left(&mut self, left: TreapTree<DataType>) {
86 if let Some(tag) = self.data.get_tag() {
87 self._right = self._right.add_tag(tag);
88 }
89 self.data.reset_tag();
90 self._left = left;
91 self.update();
92 }
93
94 fn set_right(&mut self, right: TreapTree<DataType>) {
95 if let Some(tag) = self.data.get_tag() {
96 self._left = self._left.add_tag(tag);
97 }
98 self.data.reset_tag();
99 self._right = right;
100 self.update();
101 }
102}
103
104impl<DataType: TreapNodeData> Default for TreapTree<DataType> {
105 fn default() -> Self {
106 Self(None)
107 }
108}
109
110impl<DataType: TreapNodeData> Deref for TreapTree<DataType> {
111 type Target = Option<Arc<TreapNode<DataType>>>;
112 fn deref(&self) -> &Self::Target {
113 &self.0
114 }
115}
116
117impl<DataType: TreapNodeData> TreapTree<DataType> {
118 pub fn new(data: DataType) -> Self {
119 Self(Some(Arc::new(TreapNode::new(data))))
120 }
121
122 pub fn new_from_rng<R: FnMut() -> u64>(data: DataType, rng: R) -> Self {
123 Self(Some(Arc::new(TreapNode::new_from_rng(data, rng))))
124 }
125
126 pub fn root_data_ref(&self) -> Option<&DataType> {
127 self.as_ref().map(|x| &x.data)
128 }
129
130 #[must_use]
131 pub fn map<F: FnOnce(&mut TreapNode<DataType>)>(&self, f: F) -> Self {
132 if let Some(node_ref) = self.deref() {
133 let mut node = node_ref.deref().clone();
134 f(&mut node);
135 Self(Some(Arc::new(node)))
136 } else {
137 Self::default()
138 }
139 }
140
141 #[must_use]
142 pub fn add_tag(&self, tag: DataType::TagType) -> Self {
143 self.map(|node| node.add_tag(tag))
144 }
145
146 #[must_use]
147 pub fn merge(&self, other: &Self) -> Self {
148 match (self.deref(), other.deref()) {
149 (None, None) => Self::default(),
150 (None, Some(_)) => other.clone(),
151 (Some(_), None) => self.clone(),
152 (Some(left), Some(right)) => {
153 if left.height > right.height {
154 let mut u = left.deref().to_owned();
155 u.set_right(u.get_right().merge(other));
156 Self(Some(Arc::new(u)))
157 } else {
158 let mut v = right.deref().to_owned();
159 v.set_left(self.merge(&v.get_left()));
160 Self(Some(Arc::new(v)))
161 }
162 }
163 }
164 }
165
166 #[must_use]
167 pub fn split<F: FnMut(&mut TreapNode<DataType>) -> SplitTo>(&self, mut f: F) -> (Self, Self) {
168 if let Some(node_ref) = self.deref() {
169 let mut node = node_ref.deref().clone();
170 match f(&mut node) {
171 SplitTo::Left => {
172 let (left, right) = node.get_right().split(f);
173 node.set_right(left);
174 (Self(Some(Arc::new(node))), right)
175 }
176 SplitTo::Right => {
177 let (left, right) = node.get_left().split(f);
178 node.set_left(right);
179 (left, Self(Some(Arc::new(node))))
180 }
181 }
182 } else {
183 (Self::default(), Self::default())
184 }
185 }
186
187 #[must_use]
188 pub fn query<F: FnMut(&TreapNode<DataType>) -> std::cmp::Ordering>(
189 &self,
190 mut f: F,
191 ) -> Option<Cow<'_, DataType>> {
192 if let Some(node_ref) = self.deref() {
193 match f(node_ref) {
194 std::cmp::Ordering::Equal => Some(Cow::Borrowed(&node_ref.data)),
195 std::cmp::Ordering::Less => match node_ref.get_left() {
196 Cow::Borrowed(left) => left.query(f),
197 Cow::Owned(left) => left.query(f).map(|x| Cow::Owned(x.into_owned())),
198 },
199 std::cmp::Ordering::Greater => match node_ref.get_right() {
200 Cow::Borrowed(right) => right.query(f),
201 Cow::Owned(right) => right.query(f).map(|x| Cow::Owned(x.into_owned())),
202 },
203 }
204 } else {
205 None
206 }
207 }
208
209 pub fn for_each<F: FnMut(DataType)>(&self, f: &mut F) {
210 if let Some(node_ref) = self.deref() {
211 node_ref.get_left().for_each(f);
212 f(node_ref.data.clone());
213 node_ref.get_right().for_each(f);
214 }
215 }
216
217 pub fn is_empty(&self) -> bool {
218 self.is_none()
219 }
220}