Maps vs For Loops

[stevensonmt]

I'm working through some exercises and trying different approaches. For the problem of counting primes less than a given limit, I have used two sieve approaches that I think should be equivalent. The first uses for loops plus iterators and the second avoids for loops for just map and filter methods on iterators. The for loop version is about 3 times faster and uses 3 times less memory, according to the leetcode site. Can anyone offer insight into why that is?

For loop version:

pub fn count_primes(n: i32) -> i32 {
        let upper_limit = n as usize;
        let mut prime_indices = vec![true; upper_limit];

        for p in prime_indices.iter_mut()
                             .skip(4)
                             .step_by(2) {
            *p = false;
        }

        for i in (3..upper_limit).step_by(2) {
            if prime_indices[i] {
                for p in prime_indices.iter_mut()
                                      .skip(i * i)
                                      .step_by(2 * i) {
                    *p = false;
                }
            }
        }

        prime_indices
            .into_iter()
            .enumerate()
            .skip(2)
            .filter(|&(_, p)| p)
            .map(|(i, _)| i as u64)
            .count() as i32
    }

Iterator/maps version:

pub fn count_primes(n: i32) -> i32 {
        
        let upper_limit = n as usize;
        let mut prime_indices = vec![true; upper_limit];

        prime_indices.iter_mut()
                     .skip(4)
                     .step_by(2)
                     .map(|p| *p = !*p)
                     .count(); //necessary to consume the iterator, seemed less costly than a `collect`

        let candidates = (3..upper_limit).into_iter()
                                         .step_by(2)
                                         .filter(|&i| prime_indices[i])
                                         .collect::<Vec<_>>();
        candidates.iter()
                         .map(|i| prime_indices.iter_mut()
                                               .skip(i*i)
                                               .step_by(2*i)
                                               .map(|p| *p = false)
                                               .count())
                         .count();
        
        prime_indices
            .into_iter()
            .skip(2)
            .filter(|&p| p)
            .count() as i32
    }

[jonh]

candidates is a Vec (extra space) and makes/runs iterators out of values the loop skips ( if prime_indices[i]) eg 9

[scottmcm]

@stevensonmt

.count(); //necessary to consume the iterator, seemed less costly than a collect

Sounds like you wanted for_each instead of map.

[stevensonmt]

@jonh

candidates is a Vec (extra space) and makes/runs iterators out of values the loop skips ( if prime_indices[i] ) eg 9

True, but I was creating that due to an issue with chaining filter and map where the filter borrowed prime_indices and map mutated prime_indices. Dumb solution but it worked. The better solution is to move the conditional out of a filter call and just put it in the for_each call that replaces the map. With that change the benchmarks are equivalent.

[Yandros]

Note: to force eval an Iterator<Item = ()> you can .collect::<()>() it into a (). But more generally, to force eval any Iterator you .for_each(core::mem::drop) it. And if you find yourself with .map(...).for_each(drop) you can then reduce it to .for_each(...) (with an added semi-colon)

fn count_primes (upper_limit: usize) -> usize
{
    let mut primes_sieve = vec![true; upper_limit];

    fn unset (is_prime: &'_ mut bool) { *is_prime = false; }

    // remove evens > 2
    primes_sieve
        .iter_mut()
        .skip(4)    // skip 0, 1, 2, 3
        .step_by(2) // skip odds
        .for_each(unset)
    ;

    // odds >= 3
    let odds = (3 .. upper_limit).step_by(2);

    odds.for_each(|n| if primes_sieve[n] {
        primes_sieve
            .iter_mut()
            .skip(n * n)    // for all k < n, k * n has already been sieved
            .step_by(2 * n) // (n + (2*k+1)) * n is even and has thus been sieved
            .for_each(unset)
        ;
    });
    
    // bool is isomorphic to Option<UnitType>
    type IsPrime = Option<Prime>; struct Prime;

    primes_sieve
        .into_iter()
        .skip(2)
        // equivalent to `.filter(|&is_prime| is_prime)`
        .filter_map(|is_prime: bool| -> IsPrime {
            if is_prime { Some(Prime) } else { None }
        })
        .count()
}

gives (Playground):

Counted 9592 primes under 100000 in 461.638µs
Counted 9592 primes under 100000 in 577.6µs
Counted 9592 primes under 100000 in 524.882µs