Prime sieve divs (#78)

* first draft

* finish

* revert

* consistent var name

* better way

* nits

* golf

---------

Co-authored-by: Luke Videckis <lukevideckis@gmail.com>

Author: lrvideckis

Cargo.toml

@@ -80,6 +80,10 @@ path = "examples/numbers/primes_factorize.rs"
 name = "primes_is_prime"
 path = "examples/numbers/primes_is_prime.rs"
 
+[[example]]
+name = "primes_divisorize"
+path = "examples/numbers/primes_divisorize.rs"
+
 [[example]]
 name = "manacher"
 path = "examples/strings/manacher.rs"

examples/numbers/primes_divisorize.rs

@@ -0,0 +1,24 @@
+// verification-helper: PROBLEM https://onlinejudge.u-aizu.ac.jp/courses/lesson/2/ITP1/all/ITP1_1_A
+
+use programming_team_code_rust::numbers::primes::Primes;
+
+fn main() {
+    let n = 100_000;
+    let mut divs = vec![vec![]; n + 1];
+    for i in 1..=n {
+        for j in 1..=n / i {
+            divs[i * j].push(i);
+        }
+    }
+
+    let primes = Primes::new(n + 1);
+
+    #[allow(clippy::needless_range_loop)]
+    for i in 1..=n {
+        let mut all_divisors = primes.divisorize(i);
+        all_divisors.sort();
+        assert_eq!(all_divisors, divs[i]);
+    }
+
+    println!("Hello World");
+}

examples/numbers/primes_factorize.rs

@@ -19,12 +19,12 @@ fn main() {
         let mut exps = vec![0; 1001];
         let mut prev_factor = 0;
 
-        primes.factorize(x, |factor| {
+        for factor in primes.factorize(x) {
             assert!(prev_factor <= factor);
             exps[factor] += 1;
             lcm_exps[factor] = lcm_exps[factor].max(exps[factor]);
             prev_factor = factor;
-        });
+        }
     }
 
     let mut lcm = 1;

src/numbers/primes.rs

@@ -7,9 +7,9 @@
 /// let primes = Primes::new(100);
 /// assert_eq!(primes.is_prime(2), true);
 /// assert_eq!(primes.is_prime(4), false);
-/// let mut factors = vec![];
-/// primes.factorize(12, |factor| factors.push(factor));
-/// assert_eq!(factors, vec![2, 2, 3]);
+/// assert_eq!(primes.factorize(12), [2, 2, 3]);
+/// assert_eq!(primes.divisorize(12), [1, 2, 4, 3, 6, 12]);
+/// assert!(std::panic::catch_unwind(|| primes.divisorize(0)).is_err());
 /// ```
 pub struct Primes {
     min_fact: Vec<usize>,
@@ -46,16 +46,40 @@ impl Primes {
         x >= 2 && self.min_fact[x] == x
     }
 
-    /// Calls closure on each prime factor in ascending order
+    /// Returns a vector of prime factors of the given number
+    /// The factors are sorted in ascending order
     ///
     /// # Complexity
     /// - Time: O(log x)
     /// - Space: O(log x)
-    pub fn factorize(&self, mut x: usize, mut f: impl FnMut(usize)) {
+    pub fn factorize(&self, mut x: usize) -> Vec<usize> {
+        let mut facts = vec![];
         while x > 1 {
             let p = self.min_fact[x];
-            f(p);
+            facts.push(p);
             x /= p;
         }
+        facts
+    }
+
+    /// Returns a vector of all divisors of the given number
+    ///
+    /// # Complexity
+    /// - Time: O(number_of_divisors(x))
+    /// - Space: O(number_of_divisors(x))
+    pub fn divisorize(&self, mut x: usize) -> Vec<usize> {
+        assert_ne!(x, 0);
+        let mut divs = vec![1];
+        while x > 1 {
+            let (p, len1) = (self.min_fact[x], divs.len());
+            while self.min_fact[x] == p {
+                let len2 = divs.len();
+                for i in len2 - len1..len2 {
+                    divs.push(divs[i] * p);
+                }
+                x /= p;
+            }
+        }
+        divs
     }
 }