Prime numbers (in Ruby)

For me it is like back to the basics. Back to the roots. Eratosthenes sieve. I remember when I implemented this algorithm for the first time. It was in a high school and I was feeling like doing a rocket science. Today I would like to look at the algorithm once again but from Ruby perspective.

One thing bear noting - in Ruby we have class Prime which lets you generate a sequence of prime numbers below a certain threshold. You can use it like this:

require 'prime' 
Prime.take(10) # will return iterator of 10 first prime numbers 

If you need fast and efficient solution then I recommend using Prime class. Even through I believe that my implementation that I will present in this post is correct it suffers from lack of optimization and caching.

Let's look how Eratosthenes sieve actually works. We can present algorithm in following 4 steps:

1. Create a list of consecutive integers from 2 to n (2, 3, 4, ..., n)
2. Initially, let p equal 2, the first prime number
3. Starting from p, count up in increments of p and mark each of these numbers greater than p itself in the list. These will be multiples of p: 2p, 3p, 4p, etc.; note that some of them may have already been marked.
4. Find the first number greater than p in the list that is not marked. If there was no such number, stop. Otherwise, let p now equal this number (which is the next prime), and repeat from step 3.

When the algorithm terminates, all the numbers in the list that are not marked are prime. That's all. Well, almost because I have added to my implementation two simple optimizations in step 3. The first is iterating up to square root of n. The second is that I will start marking multiplications of p from number p*p.

Now that when we know how Eratosthenes sieve works lets implement it. It took me a while before I found a way to write it in a clear and short form. After few refactoring I got to something like this:

def sieve(max=100)
    sieve = []
    (2..max).each { |i| sieve[i] = i }
    (2..Math.sqrt(max)).each do |i|
        (i*i).step(max, i) { |j| sieve[j] = nil } if sieve[i]
    end
    sieve.compact
end

It was a lot of fun writing this code ! I felt once again like a student who discovers how beautiful and simple algorithms can be.