Skip to content

Instantly share code, notes, and snippets.

@ttsugriy
Created July 30, 2023 04:09
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save ttsugriy/839a15eff4e250aac3b4ce51bf286edf to your computer and use it in GitHub Desktop.
Save ttsugriy/839a15eff4e250aac3b4ce51bf286edf to your computer and use it in GitHub Desktop.
binary_search_slice original implementation
/// Uses a sorted slice `data: &[E]` as a kind of "multi-map". The
/// `key_fn` extracts a key of type `K` from the data, and this
/// function finds the range of elements that match the key. `data`
/// must have been sorted as if by a call to `sort_by_key` for this to
/// work.
pub fn binary_search_slice<'d, E, K>(data: &'d [E], key_fn: impl Fn(&E) -> K, key: &K) -> &'d [E]
where
K: Ord,
{
let Ok(mid) = data.binary_search_by_key(key, &key_fn) else {
return &[];
};
let size = data.len();
// We get back *some* element with the given key -- so do
// a galloping search backwards to find the *first* one.
let mut start = mid;
let mut previous = mid;
let mut step = 1;
loop {
start = start.saturating_sub(step);
if start == 0 || key_fn(&data[start]) != *key {
break;
}
previous = start;
step *= 2;
}
step = previous - start;
while step > 1 {
let half = step / 2;
let mid = start + half;
if key_fn(&data[mid]) != *key {
start = mid;
}
step -= half;
}
// adjust by one if we have overshot
if start < size && key_fn(&data[start]) != *key {
start += 1;
}
// Now search forward to find the *last* one.
let mut end = mid;
let mut previous = mid;
let mut step = 1;
loop {
end = end.saturating_add(step).min(size);
if end == size || key_fn(&data[end]) != *key {
break;
}
previous = end;
step *= 2;
}
step = end - previous;
while step > 1 {
let half = step / 2;
let mid = end - half;
if key_fn(&data[mid]) != *key {
end = mid;
}
step -= half;
}
&data[start..end]
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment