nekolib/math/
factors.rs

1//! 素因数分解。
2
3/// 素因数分解。
4///
5/// $n = \\prod\_{p\_i:\\text{ prime}} p\_i^{e\_i}$ に対して、各
6/// $(p\_i, e\_i)$ を $p\_i$ の昇順に返す。
7///
8/// # Complexity
9/// $O(\\sqrt{n})$ time, $O(1)$ space.
10///
11/// # Examples
12/// ```
13/// use nekolib::math::Factors;
14///
15/// let n = 735134400_u64;
16/// let fac: Vec<_> = n.factors().collect();
17/// assert_eq!(fac, [(2, 6), (3, 3), (5, 2), (7, 1), (11, 1), (13, 1), (17, 1)]);
18/// assert_eq!(fac.iter().map(|&(p, e)| p.pow(e)).product::<u64>(), n);
19///
20/// assert_eq!(1_u64.factors().next(), None);
21/// ```
22pub trait Factors {
23    type Output;
24    fn factors(self) -> Self::Output;
25}
26
27pub struct FactorsStruct<I> {
28    i: I,
29    n: I,
30}
31
32macro_rules! impl_factors_uint {
33    ( $($ty:ty)* ) => { $(
34        impl Factors for $ty {
35            type Output = FactorsStruct<$ty>;
36            fn factors(self) -> Self::Output {
37                Self::Output { i: 2, n: self }
38            }
39        }
40        impl Iterator for FactorsStruct<$ty> {
41            type Item = ($ty, u32);
42            fn next(&mut self) -> Option<($ty, u32)> {
43                if self.n <= 1 || self.i == 0 {
44                    return None;
45                }
46                loop {
47                    match self.i.checked_pow(2) {
48                        Some(_) if self.n % self.i == 0 => {
49                            let mut e = 1;
50                            self.n /= self.i;
51                            while self.n % self.i == 0 {
52                                self.n /= self.i;
53                                e += 1;
54                            }
55                            return Some((self.i, e));
56                        }
57                        Some(ii) if ii <= self.n => {
58                            self.i += 1;
59                        }
60                        _ => {
61                            return Some((std::mem::take(&mut self.n), 1));
62                        }
63                    }
64                }
65            }
66        }
67    )* };
68}
69
70impl_factors_uint! { u8 u16 u32 u64 u128 usize }
71
72#[test]
73fn test_small() {
74    let suite: &[(u64, &[(u64, u32)])] = &[
75        (0, &[]),
76        (1, &[]),
77        (2, &[(2, 1)]),
78        (3, &[(3, 1)]),
79        (4, &[(2, 2)]),
80        (5, &[(5, 1)]),
81        (10, &[(2, 1), (5, 1)]),
82        (100, &[(2, 2), (5, 2)]),
83        (200, &[(2, 3), (5, 2)]),
84    ];
85    for (n, expected) in suite {
86        let actual: Vec<_> = n.factors().collect();
87        assert_eq!(&actual, expected);
88    }
89}
90
91#[test]
92fn test() {
93    let n = 10000_usize;
94
95    let lp = {
96        let mut lp: Vec<_> = (0..=n).collect();
97        for i in 2..=n {
98            if lp[i] < i {
99                continue;
100            }
101            for j in i..=n / i {
102                if lp[i * j] == i * j {
103                    lp[i * j] = i;
104                }
105            }
106        }
107        lp
108    };
109
110    for i in 0..=n {
111        let actual: Vec<_> = i.factors().collect();
112        let expected = {
113            let mut res = vec![];
114            let mut j = i;
115            while j > 1 {
116                if res.last().map(|&(x, _): &(usize, u32)| x) == Some(lp[j]) {
117                    res.last_mut().unwrap().1 += 1;
118                } else {
119                    res.push((lp[j], 1));
120                }
121                j /= lp[j];
122            }
123            res
124        };
125        assert_eq!(actual, expected);
126    }
127}
128
129#[test]
130fn overflow() {
131    for i in (1_u32..=1000)
132        .flat_map(|i| [i.wrapping_neg(), 2_u32.pow(16) * (2_u32.pow(16) - i)])
133    {
134        let actual: Vec<_> = i.factors().collect();
135        let expected: Vec<_> =
136            (i as u64).factors().map(|(p, e)| (p as u32, e)).collect();
137        assert_eq!(actual, expected);
138    }
139}
140
141#[test]
142fn overflow_exhaustive() {
143    for i in u8::MIN..=u8::MAX {
144        let actual: Vec<_> = i.factors().collect();
145        let expected: Vec<_> =
146            (i as u32).factors().map(|(p, e)| (p as u8, e)).collect();
147        assert_eq!(actual, expected);
148    }
149    for i in u16::MIN..=u16::MAX {
150        let actual: Vec<_> = i.factors().collect();
151        let expected: Vec<_> =
152            (i as u32).factors().map(|(p, e)| (p as u16, e)).collect();
153        assert_eq!(actual, expected);
154    }
155}
nekolib/math/factors.rs

nekolib/math/
factors.rs
