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