Skip to content

Add insert_sorted_by{,_key} methods for map/set #402

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 42 additions & 0 deletions src/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -493,6 +493,48 @@ where
}
}

/// Insert a key-value pair in the map at its ordered position among keys
/// sorted by `cmp`.
///
/// This is equivalent to finding the position with
/// [`binary_search_by`][Self::binary_search_by], then calling
/// [`insert_before`][Self::insert_before] for a new key.
///
/// If the existing keys are **not** already sorted, then the insertion
/// index is unspecified (like [`slice::binary_search`]), but the key-value
/// pair is moved to or inserted at that position regardless.
///
/// Computes in **O(n)** time (average).
pub fn insert_sorted_by<F>(&mut self, cmp: F, key: K, value: V) -> (usize, Option<V>)
where
K: Ord,
F: FnMut(&K, &V) -> Ordering,
{
let (Ok(i) | Err(i)) = self.binary_search_by(cmp);
self.insert_before(i, key, value)
}

/// Insert a key-value pair in the map at its ordered position among keys
/// sorted by `cmp`.
///
/// This is equivalent to finding the position with
/// [`binary_search_by_key`][Self::binary_search_by_key] with `cmp(key)`, then
/// calling [`insert_before`][Self::insert_before] for a new key.
///
/// If the existing keys are **not** already sorted, then the insertion
/// index is unspecified (like [`slice::binary_search`]), but the key-value
/// pair is moved to or inserted at that position regardless.
///
/// Computes in **O(n)** time (average).
pub fn insert_sorted_by_key<F, B>(&mut self, mut cmp: F, key: K, value: V) -> (usize, Option<V>)
where
B: Ord,
F: FnMut(&K, &V) -> B,
{
let (Ok(i) | Err(i)) = self.binary_search_by_key(&cmp(&key, &value), cmp);
self.insert_before(i, key, value)
}

/// Insert a key-value pair in the map before the entry at the given index, or at the end.
///
/// If an equivalent key already exists in the map: the key remains and
Expand Down
24 changes: 24 additions & 0 deletions src/map/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1226,3 +1226,27 @@ fn disjoint_indices_mut_fail_duplicate() {
Err(crate::GetDisjointMutError::OverlappingIndices)
);
}

#[test]
fn insert_sorted_by_key() {
let values = [(-1, 8), (3, 18), (-27, 2), (-2, 5)];
let mut map: IndexMap<i32, i32> = IndexMap::new();
for (key, value) in values.into_iter() {
map.insert_sorted_by_key(|k, _| k.abs(), key, value);
}
let values: Vec<_> = map.iter().map(|(k,v)|(*k,*v)).collect();
map.sort_by_cached_key(|key, _| key.abs());

assert_eq!(values, map.into_iter().collect::<Vec<_>>());
}

#[test]
fn insert_sorted_by() {
let mut values = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)];
let mut map: IndexMap<i32, i32> = IndexMap::new();
for (key, value) in values.into_iter() {
map.insert_sorted_by(|probe, _| probe.cmp(&key).reverse(), key, value);
}
values.reverse();
assert_eq!(values, map.into_iter().collect::<Vec<_>>().as_slice());
}
50 changes: 50 additions & 0 deletions src/set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -422,6 +422,56 @@ where
(index, existing.is_none())
}

/// Insert the value into the set at its ordered position among values
/// sorted by `cmp`.
///
/// This is equivalent to finding the position with
/// [`binary_search_by`][Self::binary_search_by], then calling
/// [`insert_before`][Self::insert_before].
///
/// If the sorted item is found in the set, it returns the index of that
/// existing item and `false`. Otherwise, it inserts the new item and
/// returns its sorted index and `true`.
///
/// If the existing items are **not** already sorted, then the insertion
/// index is unspecified (like [`slice::binary_search`]), but the value
/// is moved to or inserted at that position regardless.
///
/// Computes in **O(n)** time (average).
pub fn insert_sorted_by<F>(&mut self, mut cmp: F, value: T) -> (usize, bool)
where
T: Ord,
F: FnMut(&T) -> Ordering,
{
let (index, existing) = self.map.insert_sorted_by(|k, _| cmp(k), value, ());
(index, existing.is_none())
}

/// Insert the value into the set at its ordered position among values
/// sorted by `cmp`.
///
/// This is equivalent to finding the position with
/// [`binary_search_by_key`][Self::binary_search_by] with `cmp(key)`, then
/// calling [`insert_before`][Self::insert_before].
///
/// If the sorted item is found in the set, it returns the index of that
/// existing item and `false`. Otherwise, it inserts the new item and
/// returns its sorted index and `true`.
///
/// If the existing items are **not** already sorted, then the insertion
/// index is unspecified (like [`slice::binary_search`]), but the value
/// is moved to or inserted at that position regardless.
///
/// Computes in **O(n)** time (average).
pub fn insert_sorted_by_key<F, B>(&mut self, mut cmp: F, value: T) -> (usize, bool)
where
B: Ord,
F: FnMut(&T) -> B,
{
let (index, existing) = self.map.insert_sorted_by_key(|k, _| cmp(k), value, ());
(index, existing.is_none())
}

/// Insert the value into the set before the value at the given index, or at the end.
///
/// If an equivalent item already exists in the set, it returns `false` leaving the
Expand Down