diff --git a/rust/structures/treap.rs b/rust/structures/treap.rs new file mode 100644 index 00000000..80e66276 --- /dev/null +++ b/rust/structures/treap.rs @@ -0,0 +1,173 @@ +use rand::Rng; +use std::cmp::max; +use std::ops::AddAssign; + +struct Treap { + value: V, + mx: V, + add: V, + // keys should be unique + key: K, + prio: i64, + size: usize, + left: Option>>, + right: Option>>, +} + +impl Treap { + fn new(key: K, value: V) -> Option>> { + Some(Box::new(Self { + key, + value, + mx: value, + add: V::default(), + prio: rand::thread_rng().gen_range(0..1000_000_000), + size: 1, + left: None, + right: None, + })) + } + + fn apply(&mut self, v: V) { + self.value += v; + self.mx += v; + self.add += v; + } + + fn push(&mut self) { + if self.add != V::default() { + if self.left.is_some() { + self.left.as_mut().unwrap().apply(self.add); + } + if self.right.is_some() { + self.right.as_mut().unwrap().apply(self.add); + } + self.add = V::default(); + } + } + + fn pull(&mut self) { + self.size = 1 + Self::size(&self.left) + Self::size(&self.right); + self.mx = Self::mx(Self::mx(self.value, &self.left), &self.right); + } + + pub fn size(node: &Option>>) -> usize { + node.as_ref().map_or(0, |t| t.size) + } + + fn mx(value: V, node: &Option>>) -> V { + node.as_ref().map_or(value, |t| max(t.mx, value)) + } + + pub fn split( + root: Option>>, + min_right: K, + ) -> (Option>>, Option>>) { + Self::inner_split(root, min_right, |a, b| a >= b) + } + + pub fn strict_split( + root: Option>>, + min_right: K, + ) -> (Option>>, Option>>) { + Self::inner_split(root, min_right, |a, b| a > b) + } + + fn inner_split bool>( + root: Option>>, + min_right: K, + cmp: F, + ) -> (Option>>, Option>>) { + if root.is_none() { + return (None, None); + } + let mut node = root.unwrap(); + node.push(); + if cmp(node.key, min_right) { + let (l, r) = Self::inner_split(node.left, min_right, cmp); + node.left = r; + node.pull(); + (l, Some(node)) + } else { + let (l, r) = Self::inner_split(node.right, min_right, cmp); + node.right = l; + node.pull(); + (Some(node), r) + } + } + + pub fn merge( + a: Option>>, + b: Option>>, + ) -> Option>> { + if a.is_none() { + return b; + } + if b.is_none() { + return a; + } + let mut left = a.unwrap(); + let mut right = b.unwrap(); + left.push(); + right.push(); + if left.prio > right.prio { + left.right = Self::merge(left.right, Some(right)); + left.pull(); + Some(left) + } else { + right.left = Self::merge(Some(left), right.left); + right.pull(); + Some(right) + } + } + + pub fn insert(root: Option>>, key: K, value: V) -> Option>> { + let (l, r) = Self::split(root, key); + Self::merge(Self::merge(l, Treap::new(key, value)), r) + } + + pub fn remove(root: Option>>, key: K) -> Option>> { + let (l, r) = Self::split(root, key); + Self::merge(l, Self::strict_split(r, key).1) + } + + pub fn modify( + root: Option>>, + ll: K, + rr: K, + delta: V, + ) -> Option>> { + let (l1, r1) = Self::strict_split(root, rr); + let (l2, mut r2) = Self::split(l1, ll); + r2.as_mut().map(|t| t.apply(delta)); + Self::merge(Self::merge(l2, r2), r1) + } + + pub fn query(root: Option>>, ll: K, rr: K) -> (Option>>, V) { + let (l1, r1) = Self::strict_split(root, rr); + let (l2, r2) = Self::split(l1, ll); + let mx = r2.as_ref().map_or(V::default(), |t| t.mx); + (Self::merge(Self::merge(l2, r2), r1), mx) + } +} + +#[cfg(test)] +mod tests { + use crate::structures::treap::Treap; + + #[test] + fn basic_test() { + let mut treap: Option>> = None; + treap = Treap::insert(treap, 1, 10); + treap = Treap::insert(treap, 2, 20); + let (t, q1) = Treap::query(treap, 1, 2); + treap = t; + assert_eq!(q1, 20); + + treap = Treap::modify(treap, 1, 2, 100); + + let (t, q2) = Treap::query(treap, 1, 2); + treap = t; + assert_eq!(q2, 120); + } +}