Skip to main content

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