Added upper/lower_bound on slices

EbTech · EbTech · commit 0421fb0a66ef · 2020-08-16T22:08:19.000-07:00
diff --git a/src/range_query/mod.rs b/src/range_query/mod.rs
@@ -6,9 +6,25 @@ pub use dynamic_arq::{ArqView, DynamicArq};
 pub use specs::ArqSpec;
 pub use static_arq::StaticArq;
 
+/// Assuming slice is sorted, returns the minimum i for which slice[i] >= key,
+/// or slice.len() if no such i exists
+pub fn slice_lower_bound<T: Ord>(slice: &[T], key: &T) -> usize {
+    slice
+        .binary_search_by(|x| x.cmp(key).then(std::cmp::Ordering::Greater))
+        .unwrap_err()
+}
+
+/// Assuming slice is sorted, returns the minimum i for which slice[i] > key,
+/// or slice.len() if no such i exists
+pub fn slice_upper_bound<T: Ord>(slice: &[T], key: &T) -> usize {
+    slice
+        .binary_search_by(|x| x.cmp(key).then(std::cmp::Ordering::Less))
+        .unwrap_err()
+}
+
 /// A simple data structure for coordinate compression
 pub struct SparseIndex {
-    coords: Vec<i64>
+    coords: Vec<i64>,
 }
 
 impl SparseIndex {
@@ -31,6 +47,22 @@ mod test {
     use super::specs::*;
     use super::*;
 
+    #[test]
+    fn test_bounds() {
+        let mut vals = vec![16, 45, 45, 45, 82];
+
+        assert_eq!(slice_upper_bound(&vals, &44), 1);
+        assert_eq!(slice_lower_bound(&vals, &45), 1);
+        assert_eq!(slice_upper_bound(&vals, &45), 4);
+        assert_eq!(slice_lower_bound(&vals, &46), 4);
+
+        vals.dedup();
+        for (i, q) in vals.iter().enumerate() {
+            assert_eq!(slice_lower_bound(&vals, q), i);
+            assert_eq!(slice_upper_bound(&vals, q), i + 1);
+        }
+    }
+
     #[test]
     fn test_coord_compress() {
         let mut coords = vec![16, 99, 45, 18];
@@ -49,7 +81,7 @@ mod test {
         let queries = vec![(0, 10), (10, 19), (20, 29)];
         let coords = queries.iter().flat_map(|&(i, j)| vec![i, j + 1]).collect();
         let index = SparseIndex::new(coords);
-        
+
         assert_eq!(index.coords, vec![0, 10, 11, 20, 30]);
     }
 
