Improving efficiency of Rust algorithm ported from Python generator

Question

I'm learning Rust by solving ProjectEuler problems.

To this end, I am trying to port a solution to problem 88 (link) in Python that heavily relies on generators to Rust (which doesn't have generators).

A natural number, N, that can be written as the sum and product of a given set of at least two natural numbers, {a1, a2, ... , ak} is called a product-> sum number: N = a1 + a2 + ... + ak = a1 × a2 × ... × ak.

For example, 6 = 1 + 2 + 3 = 1 × 2 × 3.

For a given set of size, k, we shall call the smallest N with this property a minimal product-sum number. The minimal product-sum numbers for sets of size, k = 2, 3, 4, 5, and 6 are as follows.

k=2: 4 = 2 × 2 = 2 + 2
k=3: 6 = 1 × 2 × 3 = 1 + 2 + 3
k=4: 8 = 1 × 1 × 2 × 4 = 1 + 1 + 2 + 4
k=5: 8 = 1 × 1 × 2 × 2 × 2 = 1 + 1 + 2 + 2 + 2
k=6: 12 = 1 × 1 × 1 × 1 × 2 × 6 = 1 + 1 + 1 + 1 + 2 + 6

Hence for 2≤k≤6, the sum of all the minimal product-sum numbers is 4+6+8+12 = 30; note that 8 is only counted once in the sum.

In fact, as the complete set of minimal product-sum numbers for 2≤k≤12 is {4, 6, 8, 12, 15, 16}, the sum is 61.

What is the sum of all the minimal product-sum numbers for 2≤k≤12000?

I tried doing this by basically implementing a state machine in a recursive struct that implements the Iterator trait, but it runs many times slower than the Python version, and I can't exactly figure out why.

Flamegraph implies that a lot of time is spent in allocating memory, which sort of makes sense, but I don't understand why the Rust version is so much slower.

Can anyone explain why the Rust version is so much slower and/or how to optimise it to be as fast (or faster) than the Python version, while being more idiomatic?

Thanks in advance!

Python algorithm:

def solution88(N=12000):
    import itertools as it

    def multiplicative_partitions(n, k=None, i_min=2):
        if k is None:
            # Start from k=2 to avoid the trivial partition (n,)
            for k in it.count(2):
                x = multiplicative_partitions(n, k)
                try:
                    yield next(x)
                except StopIteration:
                    return
                yield from x
        elif k <= 0:
            return
        elif k == 1:
            yield (n,)
        elif k == 2:
            sqrt_n = int(n ** 0.5)
            for i in range(i_min, sqrt_n + 1):
                if not n % i:
                    yield (i, n // i)
        else:
            sqrt_n = int(n ** 0.5)
            for i in range(i_min, sqrt_n + 1):
                if not n % i:
                    for a in multiplicative_partitions(n // i, k - 1, i):
                        yield (i, *a)

    unprocessed = N - 2 + 1
    results = [None] * unprocessed
    n = 4  # 4 = 2 + 2 = 2 * 2
    while unprocessed > 0:
        for mp in multiplicative_partitions(n):
            if n >= sum(mp):
                k = n - sum(mp) + len(mp)
                if k <= N and results[k - 2] is None:
                    results[k - 2] = n
                    unprocessed -= 1
        n += 1
    print(set(results))
    return sum(set(results))


print(f"The answer is: {solution88()}")

(Note: I didn't come up with this Python solution, but a much slower one, initially. Credit for this one goes to user '6557' on ProjectEuler)

Rust algorithm:

use itertools::Itertools;
use std::collections::HashSet;

const LIMIT: usize = 12_000;

struct MulPar {
    n: usize,
    k: usize,
    i: usize,
    sqrt_n: usize,
    inner: Option<Box<MulPar>>,
}

impl Iterator for MulPar {
    type Item = Vec<usize>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.k == 0 || self.i > self.sqrt_n {
            None
        } else if self.k == 1 {
            self.k = 0;
            Some(vec![self.n])
        } else if self.k == 2 {
            while self.i < self.sqrt_n + 1 {
                if self.n % self.i == 0 {
                    let i = self.i;
                    self.i += 1;
                    return Some(vec![i, self.n / i]);
                }
                self.i += 1;
            }
            None
        } else if let Some(mut inner) = self.inner.take() {
            if let Some(mut next) = inner.as_mut().next() {
                next.push(self.i);
                self.inner = Some(inner);
                Some(next)
            } else {
                self.inner = None;
                self.i += 1;
                if self.i > self.sqrt_n + 1 {
                    // no more partitions to yield
                    None
                } else {
                    while self.n % self.i != 0 {
                        self.i += 1;
                    }
                    self.inner = Some(Box::new({
                        let sqrt_n = ((self.n / self.i) as f64).sqrt() as usize;
                        MulPar {
                            n: self.n / self.i,
                            k: self.k - 1,
                            i: self.i,
                            sqrt_n,
                            inner: None,
                        }
                    }));
                    self.next()
                }
            }
        } else {
            if self.i > self.sqrt_n {
                None
            } else {
                while self.n % self.i != 0 {
                    self.i += 1;
                }
                self.inner = Some(Box::new({
                    let sqrt_n = ((self.n / self.i) as f64).sqrt() as usize;
                    MulPar {
                        n: self.n / self.i,
                        k: self.k - 1,
                        i: self.i,
                        sqrt_n,
                        inner: None,
                    }
                }));
                self.next()
            }
        }
    }
}

fn main() {
    let mut unprocessed = LIMIT - 2 + 1;
    let mut n = 4;
    let mut results = vec![None; unprocessed];
    while unprocessed > 0 {
        for j in 2..(LIMIT - 1) {
            for mp in multiplicative_partitions(n, j, 2) {
                let mp_sum = mp.iter().sum::<usize>();
                if n >= mp_sum {
                    let k = n - mp_sum + mp.len();
                    if k <= LIMIT && results[k - 2].is_none() {
                        results[k - 2] = Some(n);
                        unprocessed -= 1;
                    }
                }
            }
        }
        n += 1;
    }
    let set: HashSet<usize> = HashSet::from_iter(results.into_iter().map(|x| x.unwrap_or(0)));
    println!("{:?}", set.iter().cloned().sorted().collect::<Vec<usize>>());
    let answer = set.into_iter().sum::<usize>();
    println!("The answer is: {answer}");
}

fn multiplicative_partitions(n: usize, k: usize, i: usize) -> MulPar {
    let sqrt_n = (n as f64).sqrt() as usize;
    MulPar {
        n,
        k,
        i,
        sqrt_n,
        inner: None,
    }
}

Answer 1

The fundamental reason that your original Rust code was slow was because you incorrectly translated this bit of python code:

for k in it.count(2):
    x = multiplicative_partitions(n, k)
    try:
        yield next(x)
    except StopIteration:
        return
    yield from x

Your Rust equivalent is:

for j in 2..(LIMIT - 1) {
    for mp in multiplicative_partitions(n, j, 2) {
        ...
    }
}

The key difference is that Python has an early termination criterion, it breaks out of the loop when multiplicative_partitions returns an empty sequence. But your Rust code doesn't have this, it looks all the way to LIMIT in every case which takes a long time even with Rust's speed.

Your second version was faster because you added the early breaking logic back where it was supposed to be:

let mut res = Vec::new();
for k in 2.. {
    let x = multiplicative_partitions(n, Some(k), i_min);
    if x.is_empty() {
        break;
    }
    res.extend(x.clone());
}
res

Answer 2

I ended up refactoring the code to just return Vecs from the recursive function and not mess with Iterators at all. This did the trick in terms of speeding it up to running in under 500ms, although I still don't understand why this would be so much faster.

use std::collections::{HashMap, HashSet};

const LIMIT: usize = 12_000;

fn multiplicative_partitions(n: usize, k: Option<usize>, i_min: usize) -> Vec<Vec<usize>> {
    match k {
        None => {
            let mut res = Vec::new();
            for k in 2.. {
                let x = multiplicative_partitions(n, Some(k), i_min);
                if x.is_empty() {
                    break;
                }
                res.extend(x.clone());
            }
            res
        }
        Some(k) => {
            if k == 0 {
                Vec::new()
            } else if k == 1 {
                vec![vec![n]]
            } else if k == 2 {
                let sqrt_n = (n as f64).sqrt() as usize;
                let mut res = Vec::new();
                for i in i_min..=sqrt_n {
                    if n % i == 0 {
                        res.push(vec![i, n / i]);
                    }
                }
                res
            } else {
                let sqrt_n = (n as f64).sqrt() as usize;
                let mut res = Vec::new();
                for i in i_min..=sqrt_n {
                    if n % i == 0 {
                        for mut a in multiplicative_partitions(n / i, Some(k - 1), i) {
                            a.push(i);
                            res.push(a);
                        }
                    }
                }
                res
            }
        }
    }
}

fn main() {
    let mut unprocessed = LIMIT - 2 + 1;
    let mut results = HashMap::new();
    let mut n = 4;
    while unprocessed > 0 {
        for mp in multiplicative_partitions(n, None, 2) {
            let sum = mp.iter().sum::<usize>();
            if n >= sum {
                let k = n - sum + mp.len();
                if k <= LIMIT && !results.contains_key(&k) {
                    results.insert(k, n);
                    unprocessed -= 1;
                }
            }
        }
        n += 1;
    }
    let answer: usize = results
        .into_values()
        .collect::<HashSet<usize>>()
        .into_iter()
        .sum();
    println!("The answer is: {answer}");
}

If anyone can explain why the other solutions are so much slower, I'd still love to hear it!

Answer 3

Here's how you could use internal iteration. It does have to do some allocation, but it's hopefully not as much.

use itertools::Itertools;
use std::collections::HashSet;

const LIMIT: usize = 12_000;

fn main() {
    let mut unprocessed = LIMIT - 2 + 1;
    let mut n = 4;
    let mut results = vec![None; unprocessed];
    while unprocessed > 0 {
        for j in 2..(LIMIT - 1) {
            multiplicative_partitions(n, j, 2, Box::new(|mp| {
                let mp_sum = mp.iter().sum::<usize>();
                if n >= mp_sum {
                    let k = n - mp_sum + mp.len();
                    if k <= LIMIT && results[k - 2].is_none() {
                        results[k - 2] = Some(n);
                        unprocessed -= 1;
                    }
                }
            }))
        }
        n += 1;
    }
    let set: HashSet<usize> = HashSet::from_iter(results.into_iter().map(|x| x.unwrap_or(0)));
    println!("{:?}", set.iter().cloned().sorted().collect::<Vec<usize>>());
    let answer = set.into_iter().sum::<usize>();
    println!("The answer is: {answer}");
}

fn multiplicative_partitions<'a>(n: usize, k: usize, mut i: usize, mut callback: Box<dyn FnMut(Vec<usize>) + 'a>) {
    let sqrt_n = (n as f64).sqrt() as usize;
    
    loop {
        if k == 0 || i > sqrt_n {
            break;
        } else if k == 1 {
            callback(vec![n]);
            break;
        } else if k == 2 {
            while i < sqrt_n + 1 {
                if n % i == 0 {
                    callback(vec![i, n / i]);
                }
                i += 1;
            }
            break;
        } else {
            if i > sqrt_n {
                break;
            } else {
                while i <= sqrt_n + 1 {
                    while n % i != 0 {
                        i += 1;
                    }
                    multiplicative_partitions(n / i, k - 1, i, Box::new(|mut next| {
                        next.push(i);
                        callback(next);
                    }));
                    i += 1;
                }
            }
        }
    }
}

As for my comment about using Vec, the idea would be to manually keep track of the call stack.