]>
Commit | Line | Data |
---|---|---|
f20569fa XL |
1 | #[cfg(test)] |
2 | mod tests; | |
3 | ||
4 | /// Uses a sorted slice `data: &[E]` as a kind of "multi-map". The | |
5 | /// `key_fn` extracts a key of type `K` from the data, and this | |
6 | /// function finds the range of elements that match the key. `data` | |
7 | /// must have been sorted as if by a call to `sort_by_key` for this to | |
8 | /// work. | |
9 | pub fn binary_search_slice<E, K>(data: &'d [E], key_fn: impl Fn(&E) -> K, key: &K) -> &'d [E] | |
10 | where | |
11 | K: Ord, | |
12 | { | |
13 | let mid = match data.binary_search_by_key(key, &key_fn) { | |
14 | Ok(mid) => mid, | |
15 | Err(_) => return &[], | |
16 | }; | |
17 | let size = data.len(); | |
18 | ||
19 | // We get back *some* element with the given key -- so do | |
20 | // a galloping search backwards to find the *first* one. | |
21 | let mut start = mid; | |
22 | let mut previous = mid; | |
23 | let mut step = 1; | |
24 | loop { | |
25 | start = start.saturating_sub(step); | |
26 | if start == 0 || key_fn(&data[start]) != *key { | |
27 | break; | |
28 | } | |
29 | previous = start; | |
30 | step *= 2; | |
31 | } | |
32 | step = previous - start; | |
33 | while step > 1 { | |
34 | let half = step / 2; | |
35 | let mid = start + half; | |
36 | if key_fn(&data[mid]) != *key { | |
37 | start = mid; | |
38 | } | |
39 | step -= half; | |
40 | } | |
41 | // adjust by one if we have overshot | |
42 | if start < size && key_fn(&data[start]) != *key { | |
43 | start += 1; | |
44 | } | |
45 | ||
46 | // Now search forward to find the *last* one. | |
47 | let mut end = mid; | |
48 | let mut previous = mid; | |
49 | let mut step = 1; | |
50 | loop { | |
51 | end = end.saturating_add(step).min(size); | |
52 | if end == size || key_fn(&data[end]) != *key { | |
53 | break; | |
54 | } | |
55 | previous = end; | |
56 | step *= 2; | |
57 | } | |
58 | step = end - previous; | |
59 | while step > 1 { | |
60 | let half = step / 2; | |
61 | let mid = end - half; | |
62 | if key_fn(&data[mid]) != *key { | |
63 | end = mid; | |
64 | } | |
65 | step -= half; | |
66 | } | |
67 | ||
68 | &data[start..end] | |
69 | } |