competitive/algorithm/
binary_search.rs

1use std::cmp::Ordering;
2
3/// binary search helper
4pub trait Bisect: Clone {
5    /// Return between two elements if search is not end.
6    fn middle_point(&self, other: &Self) -> Option<Self>;
7}
8
9macro_rules! impl_bisect_unsigned {
10    ($($t:ty)*) => {
11        $(impl Bisect for $t {
12            fn middle_point(&self, other: &Self) -> Option<Self> {
13                let (diff, small) = if self > other { (self - other, other) } else { (other - self, self) };
14                if diff > 1 { Some(small + diff / 2) } else { None }
15            }
16        })*
17    };
18}
19macro_rules! impl_bisect_signed {
20    ($($t:ty)*) => {
21        $(impl Bisect for $t {
22            fn middle_point(&self, other: &Self) -> Option<Self> {
23                if self.signum() != other.signum() {
24                    if match self.cmp(other) {
25                        Ordering::Less => self + 1 < *other,
26                        Ordering::Equal => false,
27                        Ordering::Greater => other + 1 < *self,
28                    } {
29                        Some((self + other) / 2)
30                    } else {
31                        None
32                    }
33                } else {
34                    let (diff, small) = if self > other { (self - other, other) } else { (other - self, self) };
35                    if diff > 1 { Some(small + diff / 2) } else { None }
36                }
37            }
38        })*
39    };
40}
41macro_rules! impl_bisect_float {
42    ($({$t:ident $u:ident $i:ident $e:expr})*) => {
43        $(impl Bisect for $t {
44            fn middle_point(&self, other: &Self) -> Option<Self> {
45                fn to_float_ord(x: $t) -> $i {
46                    let a = x.to_bits() as $i;
47                    a ^ (((a >> $e) as $u) >> 1) as $i
48                }
49                fn from_float_ord(a: $i) -> $t {
50                    $t::from_bits((a ^ (((a >> $e) as $u) >> 1) as $i) as _)
51                }
52                <$i as Bisect>::middle_point(&to_float_ord(*self), &to_float_ord(*other)).map(from_float_ord)
53            }
54        })*
55    };
56}
57impl_bisect_unsigned!(u8 u16 u32 u64 u128 usize);
58impl_bisect_signed!(i8 i16 i32 i64 i128 isize);
59impl_bisect_float!({f32 u32 i32 31} {f64 u64 i64 63});
60
61/// binary search for monotone segment
62///
63/// if `ok < err` then search [ok, err) where t(`ok`), t, t, .... t, t(`ret`), f,  ... f, f, f, `err`
64///
65/// if `err < ok` then search (err, ok] where `err`, f, f, f, ... f, t(`ret`), ... t, t, t(`ok`)
66pub fn binary_search<T, F>(mut f: F, mut ok: T, mut err: T) -> T
67where
68    T: Bisect,
69    F: FnMut(&T) -> bool,
70{
71    while let Some(m) = ok.middle_point(&err) {
72        if f(&m) {
73            ok = m;
74        } else {
75            err = m;
76        }
77    }
78    ok
79}
80
81/// binary search for slice
82pub trait SliceBisectExt<T> {
83    /// Returns the first element that satisfies a predicate.
84    fn find_bisect(&self, f: impl FnMut(&T) -> bool) -> Option<&T>;
85    /// Returns the last element that satisfies a predicate.
86    fn rfind_bisect(&self, f: impl FnMut(&T) -> bool) -> Option<&T>;
87    /// Returns the first index that satisfies a predicate.
88    /// if not found, returns `len()`.
89    fn position_bisect(&self, f: impl FnMut(&T) -> bool) -> usize;
90    /// Returns the last index+1 that satisfies a predicate.
91    /// if not found, returns `0`.
92    fn rposition_bisect(&self, f: impl FnMut(&T) -> bool) -> usize;
93}
94impl<T> SliceBisectExt<T> for [T] {
95    fn find_bisect(&self, f: impl FnMut(&T) -> bool) -> Option<&T> {
96        self.get(self.position_bisect(f))
97    }
98    fn rfind_bisect(&self, f: impl FnMut(&T) -> bool) -> Option<&T> {
99        let pos = self.rposition_bisect(f);
100        if pos == 0 { None } else { self.get(pos - 1) }
101    }
102    fn position_bisect(&self, mut f: impl FnMut(&T) -> bool) -> usize {
103        binary_search(|i| f(&self[*i as usize]), self.len() as i64, -1) as usize
104    }
105    fn rposition_bisect(&self, mut f: impl FnMut(&T) -> bool) -> usize {
106        binary_search(|i| f(&self[i - 1]), 0, self.len() + 1)
107    }
108}
109
110pub fn parallel_binary_search<T, F, G>(mut f: F, q: usize, ok: T, err: T) -> Vec<T>
111where
112    T: Bisect,
113    F: FnMut(&[Option<T>]) -> G,
114    G: Fn(usize) -> bool,
115{
116    let mut ok = vec![ok; q];
117    let mut err = vec![err; q];
118    loop {
119        let m: Vec<_> = ok
120            .iter()
121            .zip(&err)
122            .map(|(ok, err)| ok.middle_point(err))
123            .collect();
124        if m.iter().all(|m| m.is_none()) {
125            break;
126        }
127        let g = f(&m);
128        for (i, m) in m.into_iter().enumerate() {
129            if let Some(m) = m {
130                if g(i) {
131                    ok[i] = m;
132                } else {
133                    err[i] = m;
134                }
135            }
136        }
137    }
138    ok
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    const V: [i64; 10] = [0i64, 1, 1, 1, 2, 2, 3, 4, 7, 8];
145
146    #[test]
147    fn test_binary_search() {
148        assert_eq!(binary_search(|&x| V[x] >= 1, V.len(), 0), 1);
149        assert_eq!(binary_search(|&x| V[x] >= 2, V.len(), 0), 4);
150        assert_eq!(binary_search(|&x| V[x] >= 3, V.len(), 0), 6);
151        assert_eq!(binary_search(|&x| V[x] <= 1, 0, V.len()), 3);
152        assert_eq!(binary_search(|&x| V[x] <= 2, 0, V.len()), 5);
153        assert_eq!(binary_search(|&x| V[x] <= 3, 0, V.len()), 6);
154
155        assert_eq!(
156            binary_search(&|&x: &i64| V[x as usize] <= -1, -1, V.len() as i64),
157            -1
158        );
159
160        let sq2 = binary_search(|&x| x * x <= 2., 1., 4.);
161        let expect = 1.414_213_562_73;
162        assert!(expect - 1e-8 <= sq2 && sq2 <= expect + 1e-8);
163
164        assert_eq!(
165            binary_search(|&x| x < i64::MAX, i64::MIN, i64::MAX),
166            i64::MAX - 1
167        );
168        assert_eq!(
169            binary_search(|&x| x == i64::MIN, i64::MIN, i64::MAX),
170            i64::MIN
171        );
172        assert_eq!(
173            binary_search(|&x| x == i64::MAX, i64::MAX, i64::MIN),
174            i64::MAX
175        );
176        assert_eq!(
177            binary_search(|&x| x > i64::MIN, i64::MAX, i64::MIN),
178            i64::MIN + 1
179        );
180    }
181
182    #[test]
183    fn test_position() {
184        assert_eq!(V.position_bisect(|&x| x >= -1), 0);
185        assert_eq!(V.position_bisect(|&x| x >= 0), 0);
186        assert_eq!(V.position_bisect(|&x| x >= 1), 1);
187        assert_eq!(V.position_bisect(|&x| x >= 2), 4);
188        assert_eq!(V.position_bisect(|&x| x >= 3), 6);
189        assert_eq!(V.position_bisect(|&x| x >= 5), 8);
190        assert_eq!(V.position_bisect(|&x| x >= 10), 10);
191    }
192
193    #[test]
194    fn test_find() {
195        assert_eq!(V.find_bisect(|&x| x >= -1), Some(&0));
196        assert_eq!(V.find_bisect(|&x| x >= 0), Some(&0));
197        assert_eq!(V.find_bisect(|&x| x >= 1), Some(&1));
198        assert_eq!(V.find_bisect(|&x| x >= 2), Some(&2));
199        assert_eq!(V.find_bisect(|&x| x >= 3), Some(&3));
200        assert_eq!(V.find_bisect(|&x| x >= 5), Some(&7));
201        assert_eq!(V.find_bisect(|&x| x >= 10), None);
202    }
203
204    #[test]
205    fn test_rposition() {
206        assert_eq!(V.rposition_bisect(|&x| x <= -1), 0);
207        assert_eq!(V.rposition_bisect(|&x| x <= 0), 1);
208        assert_eq!(V.rposition_bisect(|&x| x <= 1), 4);
209        assert_eq!(V.rposition_bisect(|&x| x <= 2), 6);
210        assert_eq!(V.rposition_bisect(|&x| x <= 3), 7);
211        assert_eq!(V.rposition_bisect(|&x| x <= 5), 8);
212        assert_eq!(V.rposition_bisect(|&x| x <= 10), 10);
213    }
214
215    #[test]
216    fn test_rfind() {
217        assert_eq!(V.rfind_bisect(|&x| x <= -1), None);
218        assert_eq!(V.rfind_bisect(|&x| x <= 0), Some(&0));
219        assert_eq!(V.rfind_bisect(|&x| x <= 1), Some(&1));
220        assert_eq!(V.rfind_bisect(|&x| x <= 2), Some(&2));
221        assert_eq!(V.rfind_bisect(|&x| x <= 3), Some(&3));
222        assert_eq!(V.rfind_bisect(|&x| x <= 5), Some(&4));
223        assert_eq!(V.rfind_bisect(|&x| x <= 10), Some(&8));
224    }
225}
competitive/algorithm/binary_search.rs

competitive/algorithm/
binary_search.rs
