add binary_search and friends

indexmap-rs · Oct 26, 2023 · 348b42f · 348b42f
1 parent 67e67dd
commit 348b42f
Show file tree

Hide file tree

Showing 4 changed files with 185 additions and 0 deletions.
diff --git a/src/map.rs b/src/map.rs
@@ -794,6 +794,46 @@ where
         });
     }
 
+    /// Search over a sorted map with a comparator function.
+    ///
+    /// Returns the position where that value is present, or the position where can be inserted to maintain the sort.
+    /// see [slice::binary_search_by] for more details.
+    /// **O(log(n))**
+    #[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 a key extraction function.
+    ///
+    /// Returns the position where that value is present, or the position where can be inserted to maintain the sort.
+    /// see [slice::binary_search_by_key] for more details.
+    /// **O(log(n))**
+    #[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 or 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.
+    /// **O(log(n))**
+    #[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.

diff --git a/src/map/slice.rs b/src/map/slice.rs
@@ -201,6 +201,59 @@ 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 value.
+    ///
+    /// Returns the position where that value is present, or the position where can be inserted to maintain the sort.
+    /// see [slice::binary_search] for more details.
+    /// **O(log(n))**, which is notably less scalable than looking the value up in the map set is a slice from.
+    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 can be inserted to maintain the sort.
+    /// see [slice::binary_search_by] for more details.
+    /// **O(log(n))**
+    #[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 a key extraction function.
+    ///
+    /// Returns the position where that value is present, or the position where can be inserted to maintain the sort.
+    /// see [slice::binary_search_by_key] for more details.
+    /// **O(log(n))**
+    #[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 or 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.
+    /// **O(log(n))**
+    #[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> {

diff --git a/src/set.rs b/src/set.rs
@@ -688,6 +688,46 @@ where
         });
     }
 
+    /// Search over a sorted set with a comparator function.
+    ///
+    /// Returns the position where that value is present, or the position where can be inserted to maintain the sort.
+    /// see [slice::binary_search_by] for more details.
+    /// **O(log(n))**
+    #[inline]
+    pub fn binary_search_by<'a, F>(&'a self, f: F) -> Result<usize, usize>
+    where
+        F: FnMut(&'a T) -> Ordering,
+    {
+        self.as_slice().binary_search_by(f)
+    }
+
+    /// Search over a sorted set with a key extraction function.
+    ///
+    /// Returns the position where that value is present, or the position where can be inserted to maintain the sort.
+    /// see [slice::binary_search_by_key] for more details.
+    /// **O(log(n))**
+    #[inline]
+    pub fn binary_search_by_key<'a, B, F>(&'a self, b: &B, f: F) -> Result<usize, usize>
+    where
+        F: FnMut(&'a T) -> B,
+        B: Ord,
+    {
+        self.as_slice().binary_search_by_key(b, f)
+    }
+
+    /// Returns the index of the partition point or a sorted set according to the given predicate
+    /// (the index of the first element of the second partition).
+    ///
+    /// see [slice::partition_point] for more details.
+    /// **O(log(n))**
+    #[must_use]
+    pub fn partition_point<P>(&self, pred: P) -> usize
+    where
+        P: FnMut(&T) -> bool,
+    {
+        self.as_slice().partition_point(pred)
+    }
+
     /// Reverses the order of the set’s values in place.
     ///
     /// Computes in **O(n)** time and **O(1)** space.

diff --git a/src/set/slice.rs b/src/set/slice.rs
@@ -109,6 +109,58 @@ impl<T> Slice<T> {
     pub fn iter(&self) -> Iter<'_, T> {
         Iter::new(&self.entries)
     }
+
+    /// Search over a sorted set for a value.
+    ///
+    /// Returns the position where that value is present, or the position where can be inserted to maintain the sort.
+    /// see [slice::binary_search] for more details.
+    /// **O(log(n))**, which is notably less scalable than looking the value up in the set this is a slice from.
+    pub fn binary_search(&self, x: &T) -> Result<usize, usize>
+    where
+        T: Ord,
+    {
+        self.binary_search_by(|p| p.cmp(x))
+    }
+
+    /// Search over a sorted set with a comparator function.
+    ///
+    /// Returns the position where that value is present, or the position where can be inserted to maintain the sort.
+    /// see [slice::binary_search_by] for more details.
+    /// **O(log(n))**
+    #[inline]
+    pub fn binary_search_by<'a, F>(&'a self, mut f: F) -> Result<usize, usize>
+    where
+        F: FnMut(&'a T) -> Ordering,
+    {
+        self.entries.binary_search_by(move |a| f(&a.key))
+    }
+
+    /// Search over a sorted set with a key extraction function.
+    ///
+    /// Returns the position where that value is present, or the position where can be inserted to maintain the sort.
+    /// see [slice::binary_search_by_key] for more details.
+    /// **O(log(n))**
+    #[inline]
+    pub fn binary_search_by_key<'a, B, F>(&'a self, b: &B, mut f: F) -> Result<usize, usize>
+    where
+        F: FnMut(&'a T) -> B,
+        B: Ord,
+    {
+        self.binary_search_by(|k| f(k).cmp(b))
+    }
+
+    /// Returns the index of the partition point or a sorted set according to the given predicate
+    /// (the index of the first element of the second partition).
+    ///
+    /// see [slice::partition_point] for more details.
+    /// **O(log(n))**
+    #[must_use]
+    pub fn partition_point<P>(&self, mut pred: P) -> usize
+    where
+        P: FnMut(&T) -> bool,
+    {
+        self.entries.partition_point(move |a| pred(&a.key))
+    }
 }
 
 impl<'a, T> IntoIterator for &'a Slice<T> {