haar_lib/algo/
bsearch_slice.rs

1//! 二分探索
2
3macro_rules! bsearch_impl {
4    ($t:tt, $a:expr, $value:expr) => {{
5        let n = $a.len();
6        let mut b = 0;
7        let mut len = n;
8
9        while len > 0 {
10            let half = len / 2;
11            let mid = b + half;
12
13            if &$a[mid] $t $value {
14                len -= half + 1;
15                b = mid + 1;
16            } else {
17                len = half;
18            }
19        }
20
21        b
22    }}
23}
24
25/// [`BinarySearch::lower_bound`],[`BinarySearch::upper_bound`]を提供する。
26pub trait BinarySearch<T> {
27    /// x以上となる最小のindexを求める。
28    fn lower_bound(&self, value: &T) -> usize;
29    /// xを超える最小のindexを求める。
30    fn upper_bound(&self, value: &T) -> usize;
31    /// lower_bound, upper_boundの組を求める。
32    fn equal_range(&self, value: &T) -> (usize, usize) {
33        (self.lower_bound(value), self.upper_bound(value))
34    }
35}
36
37impl<T: Ord> BinarySearch<T> for [T] {
38    /// x以上となる最小のindexを求める。
39    ///
40    /// **Time complexity** $O(\log n)$
41    #[inline]
42    fn lower_bound(&self, value: &T) -> usize {
43        bsearch_impl!(<, self, value)
44    }
45
46    /// xを超える最小のindexを求める。
47    ///
48    /// **Time complexity** $O(\log n)$
49    #[inline]
50    fn upper_bound(&self, value: &T) -> usize {
51        bsearch_impl!(<=, self, value)
52    }
53}
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58    use rand::Rng;
59
60    #[test]
61    fn test_lower_bound() {
62        let a = vec![1, 1, 2, 3, 5, 6];
63        assert_eq!(a.lower_bound(&1), 0);
64        assert_eq!(a.lower_bound(&2), 2);
65        assert_eq!(a.lower_bound(&0), 0);
66        assert_eq!(a.lower_bound(&4), 4);
67        assert_eq!(a.lower_bound(&7), 6);
68    }
69
70    #[test]
71    fn test_upper_bound() {
72        let a = vec![1, 1, 2, 3, 5, 6];
73        assert_eq!(a.upper_bound(&1), 2);
74        assert_eq!(a.upper_bound(&2), 3);
75        assert_eq!(a.upper_bound(&0), 0);
76        assert_eq!(a.upper_bound(&4), 4);
77        assert_eq!(a.upper_bound(&7), 6);
78    }
79
80    #[test]
81    fn test_equal_range() {
82        let a = vec![1, 1, 3, 4, 5, 5, 5, 8];
83        assert_eq!(a.equal_range(&1), (0, 2));
84        assert_eq!(a.equal_range(&5), (4, 7));
85        assert_eq!(a.equal_range(&4), (3, 4));
86        assert_eq!(a.equal_range(&6), (7, 7));
87    }
88
89    #[test]
90    fn test_random() {
91        let mut rng = rand::thread_rng();
92
93        let n = 100;
94        let a = (0..n)
95            .map(|_| rng.gen_range(0..=10))
96            .scan(0, |state, x| {
97                *state += x;
98                Some(*state)
99            })
100            .collect::<Vec<_>>();
101
102        for x in 0..=1000 {
103            assert_eq!(
104                a.lower_bound(&x),
105                a.iter()
106                    .enumerate()
107                    .find(|(_, &y)| y >= x)
108                    .map(|(i, _)| i)
109                    .unwrap_or(n)
110            );
111
112            assert_eq!(
113                a.upper_bound(&x),
114                a.iter()
115                    .enumerate()
116                    .find(|(_, &y)| y > x)
117                    .map(|(i, _)| i)
118                    .unwrap_or(n)
119            );
120        }
121    }
122}