algorithm_rs/
search.rs

1/// Returns the smallest index on [l, r) satisfying the condition.
2pub fn lower_bound(range: std::ops::Range<usize>, prop: &dyn Fn(usize) -> bool) -> Option<usize> {
3    if prop(range.start) {
4        return Some(range.start);
5    }
6
7    let mut ng = range.start;
8    let mut ok = range.end;
9
10    while ok - ng > 1 {
11        let middle = ng + (ok - ng) / 2;
12        match prop(middle) {
13            true => ok = middle,
14            false => ng = middle,
15        }
16    }
17
18    match ok.eq(&range.end) {
19        true => None,
20        false => Some(ok),
21    }
22}
23
24#[cfg(test)]
25mod test_lower_bound {
26    use crate::search::lower_bound;
27
28    #[test]
29    fn it_works() {
30        let vs: Vec<usize> = vec![1, 2, 3, 5, 7, 10];
31        assert_eq!(lower_bound(0..vs.len(), &|x| 1 <= vs[x]), Some(0));
32        assert_eq!(lower_bound(0..vs.len(), &|x| 2 <= vs[x]), Some(1));
33        assert_eq!(lower_bound(0..vs.len(), &|x| 3 <= vs[x]), Some(2));
34        assert_eq!(lower_bound(0..vs.len(), &|x| 7 <= vs[x]), Some(4));
35        assert_eq!(lower_bound(0..vs.len(), &|x| 10 <= vs[x]), Some(5));
36        assert_eq!(lower_bound(0..vs.len(), &|x| 100 <= vs[x]), None);
37    }
38}
39
40/// Returns the largest index on [l, r) satisfying the condition.
41pub fn upper_bound(range: std::ops::Range<usize>, prop: &dyn Fn(usize) -> bool) -> Option<usize> {
42    if !prop(range.start) {
43        return None;
44    }
45
46    let mut ok = range.start;
47    let mut ng = range.end;
48
49    while ng - ok > 1 {
50        let middle = ok + (ng - ok) / 2;
51        match prop(middle) {
52            true => ok = middle,
53            false => ng = middle,
54        }
55    }
56
57    Some(ok)
58}
59
60#[cfg(test)]
61mod test_upper_bound {
62    use crate::search::upper_bound;
63
64    #[test]
65    fn it_works() {
66        let vs: Vec<usize> = vec![1, 2, 3, 5, 7, 10];
67        assert_eq!(upper_bound(0..vs.len(), &|x| vs[x] < 1), None);
68        assert_eq!(upper_bound(0..vs.len(), &|x| vs[x] < 2), Some(0));
69        assert_eq!(upper_bound(0..vs.len(), &|x| vs[x] < 3), Some(1));
70        assert_eq!(upper_bound(0..vs.len(), &|x| vs[x] < 7), Some(3));
71        assert_eq!(upper_bound(0..vs.len(), &|x| vs[x] < 10), Some(4));
72    }
73}