competitive_library/structure/
treap.rs

1#[derive(Debug, Clone)]
2pub struct Node<K, V>
3where
4    K: Ord,
5    V: Clone + Copy,
6{
7    priority: u32,
8    children: [Option<Box<Node<K, V>>>; 2],
9    key: K,
10    value: V,
11}
12impl<K, V> Node<K, V>
13where
14    K: Ord,
15    V: Clone + Copy,
16{
17    pub fn new(key: K, value: V, priority: u32) -> Self {
18        Self {
19            priority,
20            children: [None, None],
21            key,
22            value,
23        }
24    }
25    pub fn rotate(&mut self, child: usize) -> bool {
26        let a = self.children[child].take();
27        if let Some(mut node) = a {
28            std::mem::swap(self, &mut node);
29            std::mem::swap(&mut self.children[child ^ 1], &mut node.children[child]);
30            self.children[child ^ 1] = Some(node);
31            true
32        } else {
33            false
34        }
35    }
36}
37
38use crate::other::xorshift::XorShift;
39#[derive(Debug, Default)]
40pub struct Treap<K, V>
41where
42    K: Ord,
43    V: Clone + Copy,
44{
45    nodes: Option<Box<Node<K, V>>>,
46    xorshift: XorShift<u32>,
47}
48impl<K, V> Treap<K, V>
49where
50    K: Ord,
51    V: Clone + Copy,
52{
53    pub fn new() -> Self {
54        Self {
55            nodes: None,
56            xorshift: XorShift::<u32>::new(),
57        }
58    }
59
60    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
61        let new_node = Box::new(Node::new(key, value, self.xorshift.next().unwrap()));
62        Treap::insert_inner(&mut self.nodes, new_node)
63    }
64
65    fn insert_inner(node: &mut Option<Box<Node<K, V>>>, new_node: Box<Node<K, V>>) -> Option<V> {
66        if let Some(x) = node {
67            let index = match new_node.key.cmp(&x.key) {
68                std::cmp::Ordering::Equal => {
69                    let ret = std::mem::replace(&mut x.value, new_node.value);
70                    return Some(ret);
71                }
72                std::cmp::Ordering::Less => 0,
73                std::cmp::Ordering::Greater => 1,
74            };
75
76            let value = Treap::insert_inner(&mut x.children[index], new_node);
77
78            if x.priority < x.children[index].as_ref().unwrap().priority {
79                x.rotate(index);
80            }
81            value
82        } else {
83            *node = Some(new_node);
84            None
85        }
86    }
87    pub fn get(&self, key: &K) -> Option<&V> {
88        let mut node = &self.nodes;
89
90        while let Some(x) = node {
91            node = match key.cmp(&x.key) {
92                std::cmp::Ordering::Equal => return Some(&x.value),
93                std::cmp::Ordering::Less => &x.children[0],
94                std::cmp::Ordering::Greater => &x.children[1],
95            }
96        }
97        None
98    }
99
100    pub fn erase(&mut self, key: &K) -> Option<V> {
101        let mut node = &mut self.nodes;
102
103        while node.is_some() {
104            node = match key.cmp(&node.as_ref().unwrap().key) {
105                std::cmp::Ordering::Equal => {
106                    let mut y = node;
107                    loop {
108                        y = if y.as_mut().unwrap().rotate(0) {
109                            &mut y.as_mut().unwrap().children[1]
110                        } else if y.as_mut().unwrap().rotate(1) {
111                            &mut y.as_mut().unwrap().children[0]
112                        } else {
113                            return Some(y.take().unwrap().value);
114                        };
115                    }
116                }
117                std::cmp::Ordering::Less => &mut node.as_mut().unwrap().children[0],
118                std::cmp::Ordering::Greater => &mut node.as_mut().unwrap().children[1],
119            }
120        }
121
122        None
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129    #[test]
130    fn a() {
131        let mut a = Treap::new();
132        a.insert(1, 2);
133        a.insert(2, 2);
134        a.insert(3, 2);
135        a.insert(4, 2);
136        a.insert(5, 2);
137        dbg!(a);
138    }
139
140    #[test]
141    fn c() {
142        let mut a = Treap::new();
143
144        for i in 0..1000 {
145            a.insert(i, i + 10000);
146        }
147
148        assert_eq!(a.get(&0), Some(&10000));
149        assert_eq!(a.get(&5), Some(&10005));
150        assert_eq!(a.get(&10), Some(&10010));
151        assert_eq!(a.get(&100), Some(&10100));
152        assert_eq!(a.get(&999), Some(&10999));
153        assert_eq!(a.insert(999, 2), Some(10999));
154        assert_eq!(a.get(&999), Some(&2));
155    }
156
157    #[test]
158    fn heiko() {
159        let mut treap = Treap::new();
160
161        let x = XorShift::<u64>::new();
162        for i in x.take(1 << 17) {
163            treap.insert(i, ());
164        }
165
166        assert!(dbg!(f(&treap.nodes, 0)) < 100);
167    }
168
169    fn f(node: &Option<Box<Node<u64, ()>>>, count: u64) -> u64 {
170        if let Some(x) = node {
171            f(&x.children[0], count + 1).max(f(&x.children[1], count + 1))
172        } else {
173            count
174        }
175    }
176
177    #[test]
178    fn erase() {
179        let mut treap = Treap::new();
180        for i in 0..10000 {
181            treap.insert(i, i);
182        }
183
184        for i in 0..10000 {
185            if i % 2 == 0 {
186                assert_eq!(treap.erase(&i), Some(i));
187            }
188        }
189        for i in 0..10000 {
190            if i % 2 != 0 {
191                assert_eq!(treap.get(&i), Some(&i));
192            }
193        }
194    }
195}
competitive_library/structure/treap.rs

competitive_library/structure/
treap.rs
