nekolib/math/
sqrt.rs

1//! 平方根。
2
3/// 平方根。
4///
5/// $\\lfloor\\sqrt{n}\\rfloor$ を求める。
6///
7/// # Complexity
8/// $O(\\log(n))$ time.
9///
10/// Newton 法で $O(\\log(\\log(n)))$ time にするべき？
11///
12/// # Examples
13/// ```
14/// use nekolib::math::Sqrt;
15///
16/// assert_eq!(0_i32.sqrt(), 0);
17/// assert_eq!(9_i32.sqrt(), 3);
18/// assert_eq!(12_i32.sqrt(), 3);
19/// assert_eq!(16_i32.sqrt(), 4);
20///
21/// assert_eq!(u128::MAX.sqrt(), (1 << 64) - 1);
22/// ```
23///
24/// ```should_panic
25/// use nekolib::math::Sqrt;
26///
27/// (-1_i32).sqrt();
28/// ```
29pub trait Sqrt {
30    fn sqrt(self) -> Self;
31}
32
33macro_rules! impl_uint {
34    ($t:ty) => {
35        impl Sqrt for $t {
36            fn sqrt(self) -> Self {
37                let pred = |i: $t| {
38                    i.checked_pow(2).map(|i2| i2 <= self).unwrap_or(false)
39                };
40                let mut ok = 0;
41                let mut bad = 1;
42                while pred(bad) {
43                    bad *= 2;
44                }
45                while bad - ok > 1 {
46                    let mid = ok + (bad - ok) / 2;
47                    *(if pred(mid) { &mut ok } else { &mut bad }) = mid;
48                }
49                ok
50            }
51        }
52    };
53    ( $($t:ty)* ) => { $(impl_uint!($t);)* };
54}
55
56macro_rules! impl_int {
57    ($t:ty) => {
58        impl Sqrt for $t {
59            fn sqrt(self) -> Self {
60                assert!(self >= 0, "must be non-negative");
61                let pred = |i: $t| {
62                    i.checked_pow(2).map(|i2| i2 <= self).unwrap_or(false)
63                };
64                let mut ok = 0;
65                let mut bad = 1;
66                while pred(bad) {
67                    bad *= 2;
68                }
69                while bad - ok > 1 {
70                    let mid = ok + (bad - ok) / 2;
71                    *(if pred(mid) { &mut ok } else { &mut bad }) = mid;
72                }
73                ok
74            }
75        }
76    };
77    ( $($t:ty)* ) => { $(impl_int!($t);)* };
78}
79
80impl_uint!(u8 u16 u32 u64 u128 usize);
81impl_int!(i8 i16 i32 i64 i128 isize);
nekolib/math/sqrt.rs

nekolib/math/
sqrt.rs
