Skip to main content

nekolib/math/
divisors.rs

1//! 約数列挙。
2
3/// 約数列挙。
4///
5/// # Complexity
6/// $O(\\sqrt{n})$ time, $O(1)$ space.
7///
8/// # Examples
9/// ```
10/// use nekolib::math::Divisors;
11///
12/// let div: Vec<_> = 60_u64.divisors().collect();
13/// assert_eq!(div, [1, 2, 3, 4, 5, 6, 10, 12, 15, 20, 30, 60]);
14/// ```
15pub trait Divisors {
16    type Output;
17    fn divisors(self) -> Self::Output;
18}
19
20pub struct DivisorsStruct<I> {
21    below: bool,
22    i: I,
23    n: I,
24}
25
26macro_rules! impl_divisors_uint {
27    ( $($ty:ty)* ) => { $(
28        impl Divisors for $ty {
29            type Output = DivisorsStruct<$ty>;
30            fn divisors(self) -> Self::Output {
31                Self::Output { below: true, i: 0, n: self }
32            }
33        }
34        impl DivisorsStruct<$ty> {
35            fn increment(&mut self) -> Option<$ty> {
36                if !self.below {
37                    return None;
38                }
39                loop {
40                    self.i += 1;
41                    let ii = if let Some(ii) = self.i.checked_pow(2) {
42                        ii
43                    } else {
44                        self.below = false;
45                        return None;
46                    };
47                    if ii >= self.n {
48                        self.below = false;
49                        if ii > self.n {
50                            return None;
51                        }
52                    }
53                    if self.n % self.i == 0 {
54                        return Some(self.i);
55                    }
56                }
57            }
58            fn decrement(&mut self) -> Option<$ty> {
59                while self.i > 1 {
60                    self.i -= 1;
61                    if self.n % self.i == 0 {
62                        return Some(self.n / self.i);
63                    }
64                }
65                None
66            }
67        }
68        impl Iterator for DivisorsStruct<$ty> {
69            type Item = $ty;
70            fn next(&mut self) -> Option<$ty> {
71                self.increment().or_else(|| self.decrement())
72            }
73        }
74    )* };
75}
76
77impl_divisors_uint! { u8 u16 u32 u64 u128 usize }
78
79#[test]
80fn test() {
81    let n = 10000_u64;
82    for i in 0..=n {
83        let actual: Vec<_> = i.divisors().collect();
84        let expected: Vec<_> = (1..=i).filter(|&j| i % j == 0).collect();
85        assert_eq!(actual, expected);
86    }
87}
88
89#[test]
90fn overflow() {
91    for i in (1_u32..=1000)
92        .flat_map(|i| [i.wrapping_neg(), 2_u32.pow(16) * (2_u32.pow(16) - i)])
93    {
94        let actual: Vec<_> = i.divisors().collect();
95        let expected: Vec<_> =
96            (i as u64).divisors().map(|d| d as u32).collect();
97        assert_eq!(actual, expected);
98    }
99}
100
101#[test]
102fn overflow_exhaustive() {
103    for i in u8::MIN..=u8::MAX {
104        let actual: Vec<_> = i.divisors().collect();
105        let expected: Vec<_> = (i as u32).divisors().map(|d| d as u8).collect();
106        assert_eq!(actual, expected);
107    }
108    for i in u16::MIN..=u16::MAX {
109        let actual: Vec<_> = i.divisors().collect();
110        let expected: Vec<_> =
111            (i as u32).divisors().map(|d| d as u16).collect();
112        assert_eq!(actual, expected);
113    }
114}