Merge pull request #282 from Eh2406/binary_search

add binary_search and friends

Author: cuviper

src/map.rs

@@ -794,6 +794,63 @@ where
         });
     }
 
+    /// Search over a sorted map for a key.
+    ///
+    /// Returns the position where that key is present, or the position where it can be inserted to
+    /// maintain the sort. See [`slice::binary_search`] for more details.
+    ///
+    /// Computes in **O(log(n))** time, which is notably less scalable than looking the key up
+    /// using [`get_index_of`][IndexMap::get_index_of], but this can also position missing keys.
+    pub fn binary_search_keys(&self, x: &K) -> Result<usize, usize>
+    where
+        K: Ord,
+    {
+        self.as_slice().binary_search_keys(x)
+    }
+
+    /// Search over a sorted map with a comparator function.
+    ///
+    /// Returns the position where that value is present, or the position where it can be inserted
+    /// to maintain the sort. See [`slice::binary_search_by`] for more details.
+    ///
+    /// Computes in **O(log(n))** time.
+    #[inline]
+    pub fn binary_search_by<'a, F>(&'a self, f: F) -> Result<usize, usize>
+    where
+        F: FnMut(&'a K, &'a V) -> Ordering,
+    {
+        self.as_slice().binary_search_by(f)
+    }
+
+    /// Search over a sorted map with an extraction function.
+    ///
+    /// Returns the position where that value is present, or the position where it can be inserted
+    /// to maintain the sort. See [`slice::binary_search_by_key`] for more details.
+    ///
+    /// Computes in **O(log(n))** time.
+    #[inline]
+    pub fn binary_search_by_key<'a, B, F>(&'a self, b: &B, f: F) -> Result<usize, usize>
+    where
+        F: FnMut(&'a K, &'a V) -> B,
+        B: Ord,
+    {
+        self.as_slice().binary_search_by_key(b, f)
+    }
+
+    /// Returns the index of the partition point of a sorted map according to the given predicate
+    /// (the index of the first element of the second partition).
+    ///
+    /// See [`slice::partition_point`] for more details.
+    ///
+    /// Computes in **O(log(n))** time.
+    #[must_use]
+    pub fn partition_point<P>(&self, pred: P) -> usize
+    where
+        P: FnMut(&K, &V) -> bool,
+    {
+        self.as_slice().partition_point(pred)
+    }
+
     /// Reverses the order of the map’s key-value pairs in place.
     ///
     /// Computes in **O(n)** time and **O(1)** space.

src/map/slice.rs

@@ -211,6 +211,65 @@ impl<K, V> Slice<K, V> {
     pub fn into_values(self: Box<Self>) -> IntoValues<K, V> {
         IntoValues::new(self.into_entries())
     }
+
+    /// Search over a sorted map for a key.
+    ///
+    /// Returns the position where that key is present, or the position where it can be inserted to
+    /// maintain the sort. See [`slice::binary_search`] for more details.
+    ///
+    /// Computes in **O(log(n))** time, which is notably less scalable than looking the key up in
+    /// the map this is a slice from using [`IndexMap::get_index_of`], but this can also position
+    /// missing keys.
+    pub fn binary_search_keys(&self, x: &K) -> Result<usize, usize>
+    where
+        K: Ord,
+    {
+        self.binary_search_by(|p, _| p.cmp(x))
+    }
+
+    /// Search over a sorted map with a comparator function.
+    ///
+    /// Returns the position where that value is present, or the position where it can be inserted
+    /// to maintain the sort. See [`slice::binary_search_by`] for more details.
+    ///
+    /// Computes in **O(log(n))** time.
+    #[inline]
+    pub fn binary_search_by<'a, F>(&'a self, mut f: F) -> Result<usize, usize>
+    where
+        F: FnMut(&'a K, &'a V) -> Ordering,
+    {
+        self.entries.binary_search_by(move |a| f(&a.key, &a.value))
+    }
+
+    /// Search over a sorted map with an extraction function.
+    ///
+    /// Returns the position where that value is present, or the position where it can be inserted
+    /// to maintain the sort. See [`slice::binary_search_by_key`] for more details.
+    ///
+    /// Computes in **O(log(n))** time.
+    #[inline]
+    pub fn binary_search_by_key<'a, B, F>(&'a self, b: &B, mut f: F) -> Result<usize, usize>
+    where
+        F: FnMut(&'a K, &'a V) -> B,
+        B: Ord,
+    {
+        self.binary_search_by(|k, v| f(k, v).cmp(b))
+    }
+
+    /// Returns the index of the partition point of a sorted map according to the given predicate
+    /// (the index of the first element of the second partition).
+    ///
+    /// See [`slice::partition_point`] for more details.
+    ///
+    /// Computes in **O(log(n))** time.
+    #[must_use]
+    pub fn partition_point<P>(&self, mut pred: P) -> usize
+    where
+        P: FnMut(&K, &V) -> bool,
+    {
+        self.entries
+            .partition_point(move |a| pred(&a.key, &a.value))
+    }
 }
 
 impl<'a, K, V> IntoIterator for &'a Slice<K, V> {

src/map/tests.rs

@@ -447,3 +447,224 @@ fn iter_default() {
     assert_default::<ValuesMut<'static, K, V>>();
     assert_default::<IntoValues<K, V>>();
 }
+
+#[test]
+fn test_binary_search_by() {
+    // adapted from std's test for binary_search
+    let b: IndexMap<_, i32> = []
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&5)), Err(0));
+
+    let b: IndexMap<_, i32> = [4]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&3)), Err(0));
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&4)), Ok(0));
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&5)), Err(1));
+
+    let b: IndexMap<_, i32> = [1, 2, 4, 6, 8, 9]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&5)), Err(3));
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&6)), Ok(3));
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&7)), Err(4));
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&8)), Ok(4));
+
+    let b: IndexMap<_, i32> = [1, 2, 4, 5, 6, 8]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&9)), Err(6));
+
+    let b: IndexMap<_, i32> = [1, 2, 4, 6, 7, 8, 9]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&6)), Ok(3));
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&5)), Err(3));
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&8)), Ok(5));
+
+    let b: IndexMap<_, i32> = [1, 2, 4, 5, 6, 8, 9]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&7)), Err(5));
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&0)), Err(0));
+
+    let b: IndexMap<_, i32> = [1, 3, 3, 3, 7]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&0)), Err(0));
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&1)), Ok(0));
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&2)), Err(1));
+    assert!(match b.binary_search_by(|_, x| x.cmp(&3)) {
+        Ok(1..=3) => true,
+        _ => false,
+    });
+    assert!(match b.binary_search_by(|_, x| x.cmp(&3)) {
+        Ok(1..=3) => true,
+        _ => false,
+    });
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&4)), Err(4));
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&5)), Err(4));
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&6)), Err(4));
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&7)), Ok(4));
+    assert_eq!(b.binary_search_by(|_, x| x.cmp(&8)), Err(5));
+}
+
+#[test]
+fn test_binary_search_by_key() {
+    // adapted from std's test for binary_search
+    let b: IndexMap<_, i32> = []
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.binary_search_by_key(&5, |_, &x| x), Err(0));
+
+    let b: IndexMap<_, i32> = [4]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.binary_search_by_key(&3, |_, &x| x), Err(0));
+    assert_eq!(b.binary_search_by_key(&4, |_, &x| x), Ok(0));
+    assert_eq!(b.binary_search_by_key(&5, |_, &x| x), Err(1));
+
+    let b: IndexMap<_, i32> = [1, 2, 4, 6, 8, 9]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.binary_search_by_key(&5, |_, &x| x), Err(3));
+    assert_eq!(b.binary_search_by_key(&6, |_, &x| x), Ok(3));
+    assert_eq!(b.binary_search_by_key(&7, |_, &x| x), Err(4));
+    assert_eq!(b.binary_search_by_key(&8, |_, &x| x), Ok(4));
+
+    let b: IndexMap<_, i32> = [1, 2, 4, 5, 6, 8]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.binary_search_by_key(&9, |_, &x| x), Err(6));
+
+    let b: IndexMap<_, i32> = [1, 2, 4, 6, 7, 8, 9]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.binary_search_by_key(&6, |_, &x| x), Ok(3));
+    assert_eq!(b.binary_search_by_key(&5, |_, &x| x), Err(3));
+    assert_eq!(b.binary_search_by_key(&8, |_, &x| x), Ok(5));
+
+    let b: IndexMap<_, i32> = [1, 2, 4, 5, 6, 8, 9]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.binary_search_by_key(&7, |_, &x| x), Err(5));
+    assert_eq!(b.binary_search_by_key(&0, |_, &x| x), Err(0));
+
+    let b: IndexMap<_, i32> = [1, 3, 3, 3, 7]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.binary_search_by_key(&0, |_, &x| x), Err(0));
+    assert_eq!(b.binary_search_by_key(&1, |_, &x| x), Ok(0));
+    assert_eq!(b.binary_search_by_key(&2, |_, &x| x), Err(1));
+    assert!(match b.binary_search_by_key(&3, |_, &x| x) {
+        Ok(1..=3) => true,
+        _ => false,
+    });
+    assert!(match b.binary_search_by_key(&3, |_, &x| x) {
+        Ok(1..=3) => true,
+        _ => false,
+    });
+    assert_eq!(b.binary_search_by_key(&4, |_, &x| x), Err(4));
+    assert_eq!(b.binary_search_by_key(&5, |_, &x| x), Err(4));
+    assert_eq!(b.binary_search_by_key(&6, |_, &x| x), Err(4));
+    assert_eq!(b.binary_search_by_key(&7, |_, &x| x), Ok(4));
+    assert_eq!(b.binary_search_by_key(&8, |_, &x| x), Err(5));
+}
+
+#[test]
+fn test_partition_point() {
+    // adapted from std's test for partition_point
+    let b: IndexMap<_, i32> = []
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.partition_point(|_, &x| x < 5), 0);
+
+    let b: IndexMap<_, i32> = [4]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.partition_point(|_, &x| x < 3), 0);
+    assert_eq!(b.partition_point(|_, &x| x < 4), 0);
+    assert_eq!(b.partition_point(|_, &x| x < 5), 1);
+
+    let b: IndexMap<_, i32> = [1, 2, 4, 6, 8, 9]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.partition_point(|_, &x| x < 5), 3);
+    assert_eq!(b.partition_point(|_, &x| x < 6), 3);
+    assert_eq!(b.partition_point(|_, &x| x < 7), 4);
+    assert_eq!(b.partition_point(|_, &x| x < 8), 4);
+
+    let b: IndexMap<_, i32> = [1, 2, 4, 5, 6, 8]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.partition_point(|_, &x| x < 9), 6);
+
+    let b: IndexMap<_, i32> = [1, 2, 4, 6, 7, 8, 9]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.partition_point(|_, &x| x < 6), 3);
+    assert_eq!(b.partition_point(|_, &x| x < 5), 3);
+    assert_eq!(b.partition_point(|_, &x| x < 8), 5);
+
+    let b: IndexMap<_, i32> = [1, 2, 4, 5, 6, 8, 9]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.partition_point(|_, &x| x < 7), 5);
+    assert_eq!(b.partition_point(|_, &x| x < 0), 0);
+
+    let b: IndexMap<_, i32> = [1, 3, 3, 3, 7]
+        .into_iter()
+        .enumerate()
+        .map(|(i, x)| (i + 100, x))
+        .collect();
+    assert_eq!(b.partition_point(|_, &x| x < 0), 0);
+    assert_eq!(b.partition_point(|_, &x| x < 1), 0);
+    assert_eq!(b.partition_point(|_, &x| x < 2), 1);
+    assert_eq!(b.partition_point(|_, &x| x < 3), 1);
+    assert_eq!(b.partition_point(|_, &x| x < 4), 4);
+    assert_eq!(b.partition_point(|_, &x| x < 5), 4);
+    assert_eq!(b.partition_point(|_, &x| x < 6), 4);
+    assert_eq!(b.partition_point(|_, &x| x < 7), 4);
+    assert_eq!(b.partition_point(|_, &x| x < 8), 5);
+}