Browse Source

feat(bounds): implement set::lower_bound

Danilo Gómez 1 year ago
parent
commit
d2e052573d
3 changed files with 68 additions and 0 deletions
  1. 12 0
      src/node.rs
  2. 4 0
      src/set.rs
  3. 52 0
      tests/bounds.rs

+ 12 - 0
src/node.rs

@@ -132,6 +132,18 @@ impl<T: Ord> Node<T> {
             Ordering::Equal => self.count - self.left.as_ref().map_or(0, |n| n.count),
         }
     }
+
+    pub fn lower_bound(&self, value: &T) -> Option<&T> {
+        match value.cmp(&self.value) {
+            Ordering::Less => self
+                .left
+                .as_ref()
+                .and_then(|n| n.lower_bound(value))
+                .or(Some(&self.value)),
+            Ordering::Greater => self.right.as_ref().and_then(|n| n.lower_bound(value)),
+            Ordering::Equal => Some(&self.value),
+        }
+    }
 }
 
 #[cfg(test)]

+ 4 - 0
src/set.rs

@@ -44,4 +44,8 @@ impl<T: Ord> Set<T> {
     pub fn count_gt(&self, value: &T) -> i32 {
         self.count() - self.count_le(value)
     }
+
+    pub fn lower_bound(&self, value: &T) -> Option<&T> {
+        self.root.as_ref().and_then(|n| n.lower_bound(value))
+    }
 }

+ 52 - 0
tests/bounds.rs

@@ -0,0 +1,52 @@
+use mini_avl::Set;
+use std::ops::{Add, Sub};
+
+#[test]
+fn test_count() {
+    let mut set = Set::new();
+    for i in 1..11 {
+        assert!(set.insert(i));
+        assert_eq!(set.count(), i);
+    }
+    for i in 1..11 {
+        assert_eq!(set.count_lt(&i), i - 1);
+        assert_eq!(set.count_le(&i), i);
+        assert_eq!(set.count_ge(&i), 11 - i);
+        assert_eq!(set.count_gt(&i), 10 - i);
+    }
+}
+
+#[test]
+fn test_bounds_simple() {
+    let mut set = Set::new();
+    set.insert(1);
+    set.insert(5);
+    set.insert(9);
+
+    assert_eq!(set.lower_bound(&-1), Some(&1));
+    assert_eq!(set.lower_bound(&1), Some(&1));
+    assert_eq!(set.lower_bound(&3), Some(&5));
+    assert_eq!(set.lower_bound(&5), Some(&5));
+    assert_eq!(set.lower_bound(&7), Some(&9));
+    assert_eq!(set.lower_bound(&9), Some(&9));
+    assert_eq!(set.lower_bound(&11), None);
+}
+
+#[test]
+fn test_bounds() {
+    let mut values = vec![5, 3, 7, 1, 4, 6, 9];
+    let mut set = Set::new();
+    for v in values.iter() {
+        set.insert(*v);
+    }
+
+    values.sort();
+    let mut iter = values.iter();
+    let mut lb = iter.next();
+    for i in values.first().unwrap().sub(1)..values.last().unwrap().add(2) {
+        while lb.map_or(false, |v| *v < i) {
+            lb = iter.next();
+        }
+        assert_eq!(set.lower_bound(&i), lb);
+    }
+}