teardown_tree___treap/
map.rs

1use rand;
2
3use std::default::Default;
4use std::iter::{FromIterator, IntoIterator};
5use std::ops::{Index, IndexMut};
6
7use node::{Node};
8
9/// A map based on a randomized treap.
10#[derive(Debug, Clone)]
11pub struct TreapMap<K, V, Rng=rand::XorShiftRng> {
12    root: Option<Box<Node<K, V>>>,
13    size: usize,
14    rng: Rng,
15}
16
17/// An iterator over a treap's entries.
18pub struct Iter<'a, K: 'a, V: 'a> {
19    nodes: Vec<&'a Node<K, V>>,
20}
21
22/// A mutable iterator over a treap's entries.
23pub struct IterMut<'a, K: 'a, V: 'a> {
24    nodes: Vec<&'a mut Node<K, V>>,
25}
26
27/// An owning iterator over a treap's entries.
28pub struct IntoIter<K, V> {
29    nodes: Vec<Node<K, V>>,
30}
31
32enum Traversal<T> {
33    // Traverse left subtree before emitting value at node
34    Left(T),
35    // Emit value at node and continue with right subtree
36    Right(T),
37}
38
39/// An iterator over a treap's entries in key order.
40pub struct OrderedIter<'a, K: 'a, V: 'a> {
41    nodes: Vec<Traversal<&'a Node<K, V>>>,
42}
43
44impl<K: Ord, V> TreapMap<K, V, rand::XorShiftRng> {
45
46    /// Create an empty treap with the default random number generator. The
47    /// XorShift random number generator is used by default since it is fast,
48    /// but please note that it is not cryptographically secure.
49    ///
50    /// ```
51    /// let mut t = treap::TreapMap::new();
52    /// t.insert(5, "yellow");
53    /// if let Some(s) = t.get(&5) {
54    ///     println!("{}", s);
55    /// }
56    /// ```
57    pub fn new() -> TreapMap<K, V, rand::XorShiftRng> {
58        TreapMap { root: None, size: 0, rng: rand::weak_rng() }
59    }
60
61}
62
63impl<K: Ord, V, Rng: rand::Rng> TreapMap<K, V, Rng> {
64
65    /// Create an empty treap with a given random number generator.
66    ///
67    /// ```
68    /// extern crate rand;
69    ///# extern crate treap;
70    ///
71    ///# fn main() {
72    /// let mut t = treap::TreapMap::new_with_rng(rand::thread_rng());
73    /// t.insert(5, "yellow");
74    ///# }
75    /// ```
76    pub fn new_with_rng(rng: Rng) -> TreapMap<K, V, Rng> {
77        TreapMap { root: None, size: 0, rng: rng }
78    }
79
80    /// Return the number of elements in the treap.
81    ///
82    /// ```
83    /// let mut t = treap::TreapMap::new();
84    /// assert_eq!(t.len(), 0);
85    /// t.insert(5, 1);
86    /// assert_eq!(t.len(), 1);
87    /// ```
88    pub fn len(&self) -> usize { self.size }
89
90    /// Return true if the treap contains no elements.
91    ///
92    /// ```
93    /// let mut t = treap::TreapMap::new();
94    /// assert!(t.is_empty());
95    /// t.insert(5, 1);
96    /// assert!(!t.is_empty());
97    /// ```
98    pub fn is_empty(&self) -> bool { self.size == 0 }
99
100    /// Removes all elements from the treap.
101    ///
102    /// ```
103    /// let mut t = treap::TreapMap::new();
104    /// t.insert(5, 1);
105    /// t.clear();
106    /// assert!(t.is_empty());
107    /// ```
108    pub fn clear(&mut self) {
109        self.root.take();
110        self.size = 0;
111    }
112
113    /// Borrow the value corresponding to the given key if it exists in the treap.
114    ///
115    /// ```
116    /// let mut t = treap::TreapMap::new();
117    /// t.insert(5, "yellow");
118    /// t.insert(3, "blue");
119    /// t.insert(8, "green");
120    /// assert_eq!(t.get(&5), Some(&"yellow"));
121    /// assert_eq!(t.get(&10), None);
122    /// ```
123    pub fn get(&self, key: &K) -> Option<&V> {
124        self.root.as_ref().and_then(|n| n.get(key))
125    }
126
127    /// Return a mutable reference to the value corresponding to the given key if it exists in the treap.
128    ///
129    /// ```
130    /// let mut t = treap::TreapMap::new();
131    /// t.insert(5, "yellow");
132    /// match t.get_mut(&5) {
133    ///     Some(x) => *x = "blue",
134    ///     None => (),
135    /// }
136    /// assert_eq!(t.get(&5), Some(&"blue"));
137    /// ```
138    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
139        self.root.as_mut().and_then(|n| n.get_mut(key))
140    }
141
142    /// Returns true if the key is present in the treap.
143    ///
144    /// ```
145    /// let mut t = treap::TreapMap::new();
146    /// t.insert(5, "yellow");
147    /// assert_eq!(t.contains_key(&5), true);
148    /// assert_eq!(t.contains_key(&8), false);
149    /// ```
150    pub fn contains_key(&self, key: &K) -> bool {
151        self.get(key).is_some()
152    }
153
154    /// Insert a value with a given key. Returns the previous value if the key is already in the
155    /// treap.
156    ///
157    /// ```
158    /// let mut t = treap::TreapMap::new();
159    /// assert_eq!(t.insert(5, "yellow"), None);
160    /// assert_eq!(t.insert(5, "blue"), Some("yellow"));
161    /// ```
162    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
163        let priority = self.rng.next_f64();
164        let res = Node::insert_or_replace(&mut self.root, Node::new(key, value, priority));
165        if res.is_none() { self.size += 1; }
166        res
167    }
168
169    /// Remove the given key from the treap and return the value associated with it if any.
170    ///
171    /// ```
172    /// let mut t = treap::TreapMap::new();
173    /// t.insert(5, "blue");
174    /// assert_eq!(t.remove(&5), Some("blue"));
175    /// assert_eq!(t.remove(&10), None);
176    /// ```
177    pub fn remove(&mut self, key: &K) -> Option<V> {
178        let res = Node::remove(&mut self.root, key);
179        if res.is_some() { self.size -= 1; }
180        res
181    }
182
183    /// Returns an iterator over keys and values in the treap that gives the keys in sorted order.
184    ///
185    /// ```
186    /// let mut t = treap::TreapMap::new();
187    /// t.extend((1..10).map(|x| (x, "a")));
188    ///
189    /// let v: Vec<i32> = t.iter_ordered().map(|(&k, _)| k).collect();
190    /// assert_eq!(v, vec![1, 2, 3, 4, 5, 6, 7, 8, 9]);
191    /// ```
192    pub fn iter_ordered(&self) -> OrderedIter<K, V> {
193        OrderedIter {
194            nodes: match self.root {
195                None => Vec::new(),
196                Some(ref n) => vec![Traversal::Left(&**n)]
197            }
198        }
199    }
200}
201
202
203impl<K: Ord+Clone, Rng: rand::Rng> TreapMap<K, (), Rng> {
204    pub fn delete_range(&mut self, from: K, to: K, output: &mut Vec<K>) {
205        let max_prio = ::std::f64::MAX;
206        let mut root: Option<Box<Node<K, ()>>> = self.root.take();
207        let res = Node::insert_or_replace(&mut root, Node::new(from.clone(), (), max_prio));
208        let mut root = root.unwrap();
209
210        let (left, right) = (root.left.take(), root.right.take());
211        if res.is_some() {
212            output.push(root.key)
213        };
214
215        let mut root = right;
216        let res = Node::insert_or_replace(&mut root, Node::new(to.clone(), (), max_prio));
217        let mut root = root.unwrap();
218        let (mid, mut right) = (root.left.take(), root.right.take());
219        if res.is_some() {
220            let x = Node::new(to.clone(), (), self.rng.next_f64());
221            Node::insert_or_replace(&mut right, x);
222        }
223
224        *root = Node::new(from.clone(), (), max_prio);
225        root.left = left;
226        root.right = right;
227        let mut root = Some(root);
228        let res = Node::remove(&mut root, &from);
229        assert!(res.is_some());
230
231
232        let iter = IntoIter {
233            nodes: match mid {
234                None => Vec::new(),
235                Some(n) => vec![*n]
236            }
237        };
238
239        output.extend(iter.map(|(k, _)| k));
240
241
242        self.root = root;
243        self.size -= output.len();
244    }
245}
246
247
248impl<K: Ord, V, Rng: rand::Rng> Extend<(K, V)> for TreapMap<K, V, Rng> {
249    #[inline]
250    fn extend<T: IntoIterator<Item=(K, V)>>(&mut self, iter: T) {
251        for (k, v) in iter {
252            self.insert(k, v);
253        }
254    }
255}
256
257impl<K: Ord, V> FromIterator<(K, V)> for TreapMap<K, V> {
258    #[inline]
259    fn from_iter<T: IntoIterator<Item=(K, V)>>(iter: T) -> TreapMap<K, V> {
260        let mut treap = TreapMap::new();
261        treap.extend(iter);
262        treap
263    }
264}
265
266impl<K: Ord, V> Default for TreapMap<K, V> {
267    fn default() -> TreapMap<K, V> {
268        TreapMap::new()
269    }
270}
271
272/// Return an iterator that moves keys and values out of treap. The order is arbitrary.
273///
274/// ```
275/// let mut t = treap::TreapMap::new();
276/// t.extend(vec![(1, "red"), (2, "blue"), (3, "green")].into_iter());
277///
278/// // Print keys and values in arbitrary order.
279/// for (k, v) in t {
280///     println!("{}: {}", k, v);
281/// }
282/// ```
283impl<K: Ord, V, Rng: rand::Rng> IntoIterator for TreapMap<K, V, Rng> {
284    type Item = (K, V);
285    type IntoIter = IntoIter<K, V>;
286
287    fn into_iter(self) -> IntoIter<K, V> {
288        IntoIter {
289            nodes: match self.root {
290                None => Vec::new(),
291                Some(n) => vec![*n]
292            }
293        }
294    }
295}
296
297/// Return an iterator over keys and values in the treap. The order is arbitrary.
298///
299/// ```
300/// let mut t = treap::TreapMap::new();
301/// t.extend(vec![(1, 200), (2, 120), (3, 330)].into_iter());
302///
303/// let sum = (&t).into_iter().fold(0, |s, (&k, &v)| s + k + v);
304/// assert_eq!(sum, 656);
305/// ```
306impl<'a, K: Ord, V, Rng: rand::Rng> IntoIterator for &'a TreapMap<K, V, Rng> {
307    type Item = (&'a K, &'a V);
308    type IntoIter = Iter<'a, K, V>;
309
310    fn into_iter(self) -> Iter<'a, K, V> {
311        Iter {
312            nodes: match self.root {
313                None => Vec::new(),
314                Some(ref n) => vec![&**n]
315            }
316        }
317    }
318}
319
320/// Return an mutable iterator over keys and values in the treap. The order is arbitrary.
321///
322/// ```
323/// let mut t = treap::TreapMap::new();
324/// t.extend(vec![(1, 200), (2, 120), (3, 330)].into_iter());
325///
326/// for (k, v) in &mut t {
327///     *v += *k;
328/// }
329/// assert_eq!(t.get(&2), Some(&122));
330/// ```
331impl<'a, K: Ord, V, Rng: rand::Rng> IntoIterator for &'a mut TreapMap<K, V, Rng> {
332    type Item = (&'a K, &'a mut V);
333    type IntoIter = IterMut<'a, K, V>;
334
335    fn into_iter(self) -> IterMut<'a, K, V> {
336        IterMut {
337            nodes: match self.root {
338                None => Vec::new(),
339                Some(ref mut n) => vec![&mut **n]
340            }
341        }
342    }
343}
344
345impl<'a, K: Ord, V, Rng: rand::Rng> Index<&'a K> for TreapMap<K, V, Rng> {
346    type Output = V;
347
348    fn index(&self, key: &K) -> &V {
349        self.get(key).expect("no entry found for key")
350    }
351}
352
353impl<'a, K: Ord, V, Rng: rand::Rng> IndexMut<&'a K> for TreapMap<K, V, Rng> {
354    fn index_mut(&mut self, key: &K) -> &mut V {
355        self.get_mut(key).expect("no entry found for key")
356    }
357}
358
359impl<'a, K, V> Iterator for Iter<'a, K, V> {
360    type Item = (&'a K, &'a V);
361
362    fn next(&mut self) -> Option<(&'a K, &'a V)> {
363        match self.nodes.pop() {
364            None => None,
365            Some(node) => {
366                if let Some(ref boxed) = node.left {
367                    self.nodes.push(&**boxed);
368                }
369                if let Some(ref boxed) = node.right {
370                    self.nodes.push(&**boxed);
371                }
372                Some((&node.key, &node.value))
373            }
374        }
375    }
376}
377
378impl<'a, K, V> Iterator for IterMut<'a, K, V> {
379    type Item = (&'a K, &'a mut V);
380
381    fn next(&mut self) -> Option<(&'a K, &'a mut V)> {
382        match self.nodes.pop() {
383            None => None,
384            Some(node) => {
385                if let Some(boxed) = node.left.as_mut() {
386                    self.nodes.push(&mut **boxed);
387                }
388                if let Some(boxed) = node.right.as_mut() {
389                    self.nodes.push(&mut **boxed);
390                }
391                Some((&node.key, &mut node.value))
392            }
393        }
394    }
395}
396
397impl<K, V> Iterator for IntoIter<K, V> {
398    type Item = (K, V);
399
400    fn next(&mut self) -> Option<(K, V)> {
401        match self.nodes.pop() {
402            None => None,
403            Some(node) => {
404                if let Some(boxed) = node.left {
405                    self.nodes.push(*boxed);
406                }
407                if let Some(boxed) = node.right {
408                    self.nodes.push(*boxed);
409                }
410                Some((node.key, node.value))
411            }
412        }
413    }
414}
415
416impl<'a, K, V> Iterator for OrderedIter<'a, K, V> {
417    type Item = (&'a K, &'a V);
418
419    fn next(&mut self) -> Option<(&'a K, &'a V)> {
420        use self::Traversal::{Left, Right};
421        loop {
422            match self.nodes.pop() {
423                None => return None,
424                Some(Left(node)) => {
425                    self.nodes.push(Right(node));
426                    if let Some(ref node_box) = node.left {
427                        self.nodes.push(Left(&**node_box));
428                    }
429                }
430                Some(Right(node)) => {
431                    if let Some(ref node_box) = node.right {
432                        self.nodes.push(Left(&**node_box));
433                    }
434                    return Some((&node.key, &node.value));
435                }
436            }
437        }
438    }
439}
440
441#[cfg(test)]
442mod tests {
443    use super::TreapMap;
444
445    #[test]
446    fn test_len() {
447        let mut t = TreapMap::new();
448        assert_eq!(t.len(), 0);
449        t.insert(1, 1);
450        assert_eq!(t.len(), 1);
451        t.insert(1, 2);
452        assert_eq!(t.len(), 1);
453        t.insert(2, 2);
454        t.insert(3, 3);
455        assert_eq!(t.len(), 3);
456        t.remove(&2);
457        assert_eq!(t.len(), 2);
458    }
459}
teardown_tree___treap/map.rs

teardown_tree___treap/
map.rs
