#![allow(dead_code)]
#[derive(Debug)]
struct TreapNode<K, V> {
key: K,
val: V,
priority: u64,
left: Option<Box<TreapNode<K, V>>>,
right: Option<Box<TreapNode<K, V>>>,
}
impl<K: Ord, V> TreapNode<K, V> {
fn new(key: K, val: V, priority: u64) -> Box<Self> {
Box::new(TreapNode {
key,
val,
priority,
left: None,
right: None,
})
}
fn rotate_right(mut y: Box<TreapNode<K, V>>) -> Box<TreapNode<K, V>> {
let Some(mut x) = y.left.take() else { return y };
y.left = x.right.take();
x.right = Some(y);
x
}
fn rotate_left(mut x: Box<TreapNode<K, V>>) -> Box<TreapNode<K, V>> {
let Some(mut y) = x.right.take() else {
return x;
};
x.right = y.left.take();
y.left = Some(x);
y
}
fn insert(
node: Option<Box<TreapNode<K, V>>>,
key: K,
val: V,
prio: u64,
) -> (Box<TreapNode<K, V>>, bool) {
let Some(mut n) = node else {
return (Self::new(key, val, prio), true);
};
let added;
match key.cmp(&n.key) {
std::cmp::Ordering::Equal => {
n.val = val;
return (n, false);
}
std::cmp::Ordering::Less => {
let (l, a) = Self::insert(n.left.take(), key, val, prio);
n.left = Some(l);
added = a;
if n.left.as_ref().is_some_and(|l| l.priority > n.priority) {
n = Self::rotate_right(n);
}
}
std::cmp::Ordering::Greater => {
let (r, a) = Self::insert(n.right.take(), key, val, prio);
n.right = Some(r);
added = a;
if n.right.as_ref().is_some_and(|r| r.priority > n.priority) {
n = Self::rotate_left(n);
}
}
}
(n, added)
}
fn get(&self, key: &K) -> Option<&V> {
match key.cmp(&self.key) {
std::cmp::Ordering::Equal => Some(&self.val),
std::cmp::Ordering::Less => self.left.as_ref()?.get(key),
std::cmp::Ordering::Greater => self.right.as_ref()?.get(key),
}
}
}
pub struct TreapMap<K: Ord, V> {
root: Option<Box<TreapNode<K, V>>>,
count: usize,
rng: u64,
}
impl<K: Ord, V> TreapMap<K, V> {
pub fn new() -> Self {
TreapMap {
root: None,
count: 0,
rng: 0x853c_49e6_748f_ea9b,
}
}
fn next_prio(&mut self) -> u64 {
self.rng ^= self.rng << 13;
self.rng ^= self.rng >> 7;
self.rng ^= self.rng << 17;
self.rng
}
pub fn insert(&mut self, key: K, val: V) {
let prio = self.next_prio();
let (node, added) = TreapNode::insert(self.root.take(), key, val, prio);
self.root = Some(node);
if added {
self.count += 1;
}
}
pub fn get(&self, key: &K) -> Option<&V> {
self.root.as_ref()?.get(key)
}
pub fn len(&self) -> usize {
self.count
}
pub fn is_empty(&self) -> bool {
self.count == 0
}
pub fn contains_key(&self, key: &K) -> bool {
self.get(key).is_some()
}
}
impl<K: Ord, V> Default for TreapMap<K, V> {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_insert_and_get() {
let mut t: TreapMap<u32, u32> = TreapMap::new();
t.insert(5, 50);
assert_eq!(t.get(&5), Some(&50) );
}
#[test]
fn test_missing_key() {
let t: TreapMap<u32, u32> = TreapMap::new();
assert_eq!(t.get(&1), None );
}
#[test]
fn test_update() {
let mut t: TreapMap<u32, u32> = TreapMap::new();
t.insert(2, 20);
t.insert(2, 99);
assert_eq!(t.get(&2), Some(&99) );
}
#[test]
fn test_len() {
let mut t: TreapMap<u32, u32> = TreapMap::new();
t.insert(1, 1);
t.insert(2, 2);
assert_eq!(t.len(), 2);
}
#[test]
fn test_is_empty() {
let t: TreapMap<u32, u32> = TreapMap::new();
assert!(t.is_empty());
}
#[test]
fn test_many_inserts() {
let mut t: TreapMap<u32, u32> = TreapMap::new();
for i in 0u32..50 {
t.insert(i, i * 3);
}
for i in 0u32..50 {
assert_eq!(t.get(&i), Some(&(i * 3)));
}
}
#[test]
fn test_contains_key() {
let mut t: TreapMap<u32, u32> = TreapMap::new();
t.insert(9, 90);
assert!(t.contains_key(&9));
assert!(!t.contains_key(&10));
}
#[test]
fn test_reverse_inserts() {
let mut t: TreapMap<u32, u32> = TreapMap::new();
for i in (0u32..30).rev() {
t.insert(i, i);
}
assert_eq!(t.len(), 30);
}
#[test]
fn test_default() {
let t: TreapMap<u32, u32> = TreapMap::default();
assert!(t.is_empty());
}
}