bisect/
lib.rs

1use std::ops::{Range, RangeFrom, RangeTo};
2
3pub trait Bisect {
4    type Input;
5    type Output;
6    fn bisect(&self, pred: impl FnMut(&Self::Input) -> bool) -> Self::Output;
7}
8
9macro_rules! impl_bisect_uint {
10    ( $($ty:ty)* ) => { $(
11        impl Bisect for Range<$ty> {
12            type Input = $ty;
13            type Output = $ty;
14            fn bisect(&self, mut pred: impl FnMut(&$ty) -> bool) -> $ty {
15                let Range { start: mut ok, end: mut bad } = *self;
16                if !pred(&ok) {
17                    return ok;
18                }
19                while bad - ok > 1 {
20                    let mid = ok + (bad - ok) / 2;
21                    *(if pred(&mid) { &mut ok } else { &mut bad }) = mid;
22                }
23                bad
24            }
25        }
26        impl Bisect for RangeFrom<$ty> {
27            type Input = $ty;
28            type Output = $ty;
29            fn bisect(&self, mut pred: impl FnMut(&$ty) -> bool) -> $ty {
30                let RangeFrom { start: ok } = *self;
31                if !pred(&ok) {
32                    return ok;
33                }
34                let mut w = 1;
35                while pred(&(ok + w)) {
36                    w *= 2;
37                }
38                (ok + w / 2..ok + w).bisect(pred)
39            }
40        }
41        impl Bisect for RangeTo<$ty> {
42            type Input = $ty;
43            type Output = $ty;
44            fn bisect(&self, mut pred: impl FnMut(&$ty) -> bool) -> $ty {
45                let RangeTo { end: bad } = *self;
46                if pred(&bad) {
47                    return bad;
48                }
49                let mut w = 1;
50                while !pred(&(bad - w)) {
51                    w *= 2;
52                }
53                (bad - w..bad - w / 2).bisect(pred)
54            }
55        }
56    )* }
57}
58
59impl_bisect_uint! { u8 u16 u32 u64 u128 usize }
60
61macro_rules! impl_bisect_int {
62    ( $( ($ity:ty, $uty:ty, $i2u:ident, $u2i:ident), )* ) => { $(
63        impl Bisect for Range<$ity> {
64            type Input = $ity;
65            type Output = $ity;
66            fn bisect(&self, mut pred: impl FnMut(&$ity) -> bool) -> $ity {
67                let Range { start, end } = *self;
68                let start = $i2u(start);
69                let end = $i2u(end);
70                $u2i((start..end).bisect(|&u| pred(&$u2i(u))))
71            }
72        }
73        impl Bisect for RangeFrom<$ity> {
74            type Input = $ity;
75            type Output = $ity;
76            fn bisect(&self, mut pred: impl FnMut(&$ity) -> bool) -> $ity {
77                let RangeFrom { start } = *self;
78                let start = $i2u(start);
79                $u2i((start..).bisect(|&u| pred(&$u2i(u))))
80            }
81        }
82        impl Bisect for RangeTo<$ity> {
83            type Input = $ity;
84            type Output = $ity;
85            fn bisect(&self, mut pred: impl FnMut(&$ity) -> bool) -> $ity {
86                let RangeTo { end } = *self;
87                let end = $i2u(end);
88                $u2i((..end).bisect(|&u| pred(&$u2i(u))))
89            }
90        }
91        fn $i2u(i: $ity) -> $uty { !(!(0 as $uty) >> 1) ^ i as $uty }
92        fn $u2i(u: $uty) -> $ity { (!(!0 >> 1) ^ u) as $ity }
93    )* }
94}
95
96impl_bisect_int! {
97    (i8, u8, i2u8, u2i8),
98    (i16, u16, i2u16, u2i16),
99    (i32, u32, i2u32, u2i32),
100    (i64, u64, i2u64, u2i64),
101    (i128, u128, i2u128, u2i128),
102    (isize, usize, i2usize, u2isize),
103}
104
105macro_rules! impl_bisect_float {
106    (
107        $(
108            (
109                $fty:ty, $ity:ty, $uty:ty, $w:literal,
110                $f2u:ident, $u2f:ident, $mask:ident
111            ),
112        )*
113    ) => { $(
114        impl Bisect for Range<$fty> {
115            type Input = $fty;
116            type Output = $fty;
117            fn bisect(&self, mut pred: impl FnMut(&$fty) -> bool) -> $fty {
118                let Range { start, end } = *self;
119                let start = $f2u(start);
120                let end = $f2u(end);
121                $u2f((start..end).bisect(|&u| pred(&$u2f(u))))
122            }
123        }
124        impl Bisect for RangeFrom<$fty> {
125            type Input = $fty;
126            type Output = $fty;
127            fn bisect(&self, mut pred: impl FnMut(&$fty) -> bool) -> $fty {
128                let RangeFrom { start } = *self;
129                let start = $f2u(start);
130                $u2f((start..).bisect(|&u| pred(&$u2f(u))))
131            }
132        }
133        impl Bisect for RangeTo<$fty> {
134            type Input = $fty;
135            type Output = $fty;
136            fn bisect(&self, mut pred: impl FnMut(&$fty) -> bool) -> $fty {
137                let RangeTo { end } = *self;
138                let end = $f2u(end);
139                $u2f((..end).bisect(|&u| pred(&$u2f(u))))
140            }
141        }
142        fn $mask(u: $uty) -> $uty {
143            ((u as $ity >> ($w - 1)) as $uty >> 1) | 1 << ($w - 1)
144        }
145        fn $f2u(f: $fty) -> $uty { f.to_bits() ^ $mask(f.to_bits()) }
146        fn $u2f(u: $uty) -> $fty { <$fty>::from_bits(u ^ $mask(!u)) }
147    )* };
148}
149
150impl_bisect_float! {
151    (f32, i32, u32, 32, f2u32, u2f32, mask32),
152    (f64, i64, u64, 64, f2u64, u2f64, mask64),
153}
154
155impl<T> Bisect for [T] {
156    type Input = T;
157    type Output = usize;
158    fn bisect(&self, mut pred: impl FnMut(&T) -> bool) -> usize {
159        if self.is_empty() || !pred(&self[0]) {
160            return 0;
161        }
162        let mut ok = 0;
163        let mut bad = self.len();
164        while bad - ok > 1 {
165            let mid = ok + (bad - ok) / 2;
166            *(if pred(&self[mid]) { &mut ok } else { &mut bad }) = mid;
167        }
168        bad
169    }
170}
171
172#[test]
173fn sanity_check() {
174    {
175        let pred = |&x: &i64| x < 100;
176        assert_eq!((0_i64..200).bisect(pred), 100);
177        assert_eq!((0_i64..).bisect(pred), 100);
178        assert_eq!((..200_i64).bisect(pred), 100);
179    }
180
181    {
182        let pred = |&x: &i64| x.abs() < 100;
183        assert_eq!((0_i64..200).bisect(pred), 100);
184        assert_eq!((0_i64..).bisect(pred), 100);
185        assert_eq!((..0_i64).bisect(|x| !pred(x)), -99);
186    }
187
188    {
189        let pred = |&x: &i64| x < 5;
190        let a = vec![0, 1, 4, 5, 5, 9];
191        assert_eq!(a.bisect(pred), 3);
192        assert_eq!(a[4..].bisect(pred), 0);
193    }
194
195    {
196        let pred = |&x: &f64| 2.0_f64.powf(x) < 3.0;
197        assert!(((1.0_f64..2.0).bisect(pred) - 3.0_f64.log2()) <= 5.0e-324);
198        // println!("{:.40}", 3.0_f64.log2());
199        // println!("{:.40}", (1.0_f64..2.0).bisect(pred));
200    }
201    {
202        assert_eq!([0, 1, 4, 5, 9].bisect(|&x: &i32| x < 5), 3);
203        assert_eq!((0..100).bisect(|&x: &i32| x * x < 200), 15);
204        assert_eq!((0..).bisect(|&x: &i32| x * x < 200), 15);
205    }
206    {
207        let pred = |&x: &i32| x * x < 200;
208        assert_eq!((0..100).bisect(pred), 15);
209        assert_eq!((0..).bisect(pred), 15);
210    }
211}
bisect/lib.rs

bisect/
lib.rs
