diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 13353021..04bd3e14 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -6,6 +6,7 @@ edition = "2021" [dev-dependencies] rstest="0.22.0" +rand = "0.8.5" [lib] path = "lib.rs" diff --git a/rust/structures/mergeable_heap.rs b/rust/structures/mergeable_heap.rs new file mode 100644 index 00000000..e50b32be --- /dev/null +++ b/rust/structures/mergeable_heap.rs @@ -0,0 +1,75 @@ +use std::mem::swap; +use std::rc::Rc; + +struct Heap { + value: V, + left: Option>>, + right: Option>>, +} + +impl Heap { + fn new(value: V) -> Option>> { + Some(Rc::new(Self { + value, + left: None, + right: None, + })) + } + fn merge(a: &Option>>, b: &Option>>) -> Option>> { + if a.is_none() { + return b.clone(); + } + if b.is_none() { + return a.clone(); + } + let mut ra = a.clone()?; + let mut rb = b.clone()?; + if ra.value > rb.value { + swap(&mut ra, &mut rb); + } + if rand::random() { + ra = Rc::new(Heap { + value: ra.value.clone(), + left: ra.right.clone(), + right: ra.left.clone(), + }); + } + Some(Rc::new(Heap { + value: ra.value.clone(), + left: Self::merge(&ra.left.clone(), &Some(rb)), + right: ra.right.clone(), + })) + } + + fn remove_min(heap: &Option>>) -> (Option>>, V) { + let h = heap.as_ref().unwrap(); + ( + Self::merge(&h.as_ref().left, &h.as_ref().right), + h.as_ref().value.clone(), + ) + } + + fn add(heap: &Option>>, value: V) -> Option>> { + Self::merge(heap, &Heap::new(value)) + } +} + +#[cfg(test)] +mod tests { + use crate::structures::mergeable_heap::Heap; + + #[test] + fn basic_test() { + let mut h = None; + h = Heap::add(&h, 3); + h = Heap::add(&h, 1); + h = Heap::add(&h, 2); + let mut values = Vec::new(); + while h.is_some() { + let (heap, min_value) = Heap::remove_min(&h); + values.push(min_value); + h = heap; + } + assert_eq!(values, [1, 2, 3]); + } +} diff --git a/rust/structures/mod.rs b/rust/structures/mod.rs index da7b429a..bdb69099 100644 --- a/rust/structures/mod.rs +++ b/rust/structures/mod.rs @@ -1,3 +1,4 @@ pub mod disjoint_sets; pub mod fenwick_tree; +pub mod mergeable_heap; pub mod persistent_tree;