Find the nth prime number

Question

It's a long time since I've written a Sieve of Eratosthenes (about 3 decades, I think, in ZX Basic). So I decided to revisit in the light of experience.

Rather than allocating a vector of unknown size to sieve values, I keep an ordered set of upcoming composite values which are updated as they are reached.

#include <map>

unsigned long find_nth_prime(unsigned long n)
{
    if (!n) return 1;           // "0th prime"
    if (!--n) return 2;         // first prime

    // The map m contains one entry for each prime.  The key is the
    // next time to update it, and the value is the amount to increment
    // by each time it is reached.
    auto m = std::map<unsigned long, unsigned long>();
    // The overflow tests below are valid only for unsigned types.
    for (auto i = 3ul;  true;  i += 2) {
        if (i < 3ul)            // overflow test
            throw std::overflow_error("prime too large");

        const auto it = m.cbegin();
        if (it == m.end() || it->first != i) {
            // it's a prime
            if (!--n)
                return i;
            if (i*i > i)        // overflow test
                m.insert({i*i, 2*i});
        } else {
            // it's composite - advance the factor
            auto next = it->first;
            do {
                next += it->second;
            } while (next > i && !m.emplace(next, it->second).second);
            m.erase(it);
        }
    }
}

And the test program (not really part of the review):

#include <iostream>
#include <string>

int main(int, char **argv)
{
    std::cout.imbue(std::locale(""));
    while (*++argv) {
        try {
            std::cout << find_nth_prime(std::stoul(*argv)) << std::endl;
        } catch (std::exception&) {
            std::cerr << "Invalid argument: " << *argv << std::endl;
        }
    }
}

Compiled with GCC:

g++ -std=c++17 -O3 -Wall -Wextra -Wpedantic

Whilst the performance is acceptable (about 1.6 seconds on my hardware to find the one-millionth prime - 15,485,863), I'm unsure whether a std::map is really the best structure for keeping track. As usual, any other critique of the code is welcome.

Answer 1

Bug: Only works if longs are 64-bit

I ran your program on my 32-bit machine and got the wrong answer for the 1 millionth prime: 15498907 instead of 15485863. The problem is here:

        if (i*i > i)        // overflow test

If longs are 32 bits, then a number such as 0x10001 will pass this overflow test (because 0x10001 * 0x10001 = 0x20001 after overflowing). This eventually causes problems because overflowed values are inserted into the map. Actually, even if longs are 64-bit, you will still have a problem, but it will only happen after i reaches 0x100000001.

I fixed the problem by doing this:

    unsigned long max_i = std::sqrt(~0UL);
    // ...
        if (i <= max_i)      // overflow test

Slow compared to vector based solution

I found your program to be a lot slower than a traditional vector based sieve. Here are timings for your program vs a traditional sieve that I wrote which used a vector<bool>:

Map    based,  1 millionth prime:  3.77 seconds
Map    based, 10 millionth prime: 47.91 seconds
Vector based,  1 millionth prime:  0.04 seconds
Vector based, 10 millionth prime:  0.93 seconds

Sample vector based code

Since there was a comment asking about how the vector based code would work, here is the full program I used:

#include <iostream>
#include <string>
#include <cstdlib>
#include <cmath>
#include <cstdio>
#include <vector>

// This function uses the prime counting function approximation:
//
// n = x / ln(x)
//
// Where n is the number of primes below the number x.  Since we are trying
// to find the nth prime, we should solve for x here to determine an
// approximate value for the nth prime.  It turns out that for n >= 6, the
// value we solve for x will always be larger than the nth prime value.
//
// If the sieve size doesn't fit in an unsigned long, this function returns 0.
unsigned long findSieveSize(unsigned long n)
{
    // For small n, the formula returns a value too low, so we can just
    // hardcode the sieve size to 5 (5th prime is 11).
    if (n < 6)
        return 13;

    // We can't find a prime that will exceed ~0UL.
    if (n >= (~0UL / std::log(~0UL)))
        return 0;

    // Binary search for the right value.
    unsigned long low  = n;
    unsigned long high = ~0UL - 1;
    do {
        unsigned long mid   = low + (high - low) / 2;
        double        guess = mid / std::log(mid);

        if (guess > n)
            high = (unsigned long) mid - 1;
        else
            low = (unsigned long) mid + 1;
    } while (low < high);
    return high + 1;
}

unsigned long find_nth_prime(unsigned long n)
{
    if (!n) return 1;           // "0th prime"
    if (!--n) return 2;         // first prime

    unsigned long sieveSize = findSieveSize(n);
    unsigned long count     = 0;
    unsigned long max_i     = std::sqrt(sieveSize-1)+1;

    if (sieveSize == 0)
        return 0;

    std::vector<bool> sieve(sieveSize);
    for (unsigned long i = 3;  true;  i += 2) {
        if (!sieve[i]) {
            if (++count == n)
                return i;
            if (i >= max_i)
                continue;
            unsigned long j    = i*i;
            unsigned long inc  = i+i;
            unsigned long maxj = sieveSize - inc;
            // This loop checks j before adding inc so that we can stop
            // before j overflows.
            do {
                sieve[j] = true;
                if (j >= maxj)
                    break;
                j += inc;
            } while (1);
        }
    }
    return 0;
}

int main(int, char **argv)
{
    std::cout.imbue(std::locale(""));
    while (*++argv) {
        try {
            std::cout << find_nth_prime(atoi(*argv)) << std::endl;
        } catch (std::exception&) {
            std::cerr << "Invalid argument: " << *argv << std::endl;
        }
    }
}

Answer 2

The treatment of the invalid argument can be improved:

if (!n) return 1;           // "0th prime"

Instead of returning a dummy value, it would be clearer to throw an exception:

if (!n) throw std::invalid_argument("n must be at least 1");

---

Also, the exception detail is lost when reporting:

std::cerr << "Invalid argument: " << *argv << std::endl;

We can print it if we name the caught exception:

std::cerr << "Invalid argument " << *argv << ": " << e.what() << std::endl;

Answer 3

Your code should express intent. Neither of these express intent. They look more like optimizations.

if (!n) return 1;           // "0th prime"
if (!--n) return 2;         // first prime

This expresses intent.

if (n == 0) return 1;           // "0th prime"
if (n == 1) return 2;           // first prime

Is this valid?

int main(int, char **argv)

This is valid:

int main(int argc, char *argv[])

Not sure this does what you think:

std::stoul(*argv)

Is argv[0] not the name of the application and argv[1] is the first user supplied argument.

Must admit I don't understand your implementation. So a decent description in comments about how that implements a sieve would definitely not be out of place.

Answer 4

Use a heap instead of a map

Standard Library provides std::priority_queue, which provides fast access to the highest value (which we can make the lowest with a suitable < operator). I was able to reduce runtime by using this class instead of std::map.

Use a wheel to eliminate many composite numbers

A 2,3,5 wheel reduced runtime by a further 50% . A large part of this gain is by skipping over wheel-composite numbers when we increment positions; this greatly reduces updates to the container.

---

Improved code

On my system, the modified program takes 0.6 seconds to find the millionth prime, compared with 1.6 seconds for the original (both using gcc -O3):

// Include the headers we'll need, and define the result type:

#include <algorithm>
#include <cmath>
#include <queue>
#include <stdexcept>
#include <tuple>

using Number = unsigned long;

// The wheel is just a table lookup:

static constexpr bool wheel_prime(Number i)
{
    // mod-30 wheel eliminates multiples of 3 and 5
    // incrementing by 2 eliminates even numbers
    constexpr int wheel_size = 30;
    constexpr bool wheel[2][wheel_size] = {
        {
            false,
            false,
            true, //  2
            true, //  3
            false,
            true, //  5
            false,
            true, //  7
            false,
            false,
            false,
            true, //  11
            false,
            true, //  13
            false,
            false,
            false,
            true, //  17
            false,
            true, //  19
            false,
            false,
            false,
            true, //  23
            false,
            false,
            false,
            false,
            false,
            true, //  29
        },
        {
            false,
            true, //  1
            false,
            false,
            false,
            false,
            false,
            true, //  7
            false,
            false,
            false,
            true, //  11
            false,
            true, //  13
            false,
            false,
            false,
            true, //  17
            false,
            true, //  19
            false,
            false,
            false,
            true, //  23
            false,
            false,
            false,
            false,
            false,
            true, //  29
        }
    };
    return wheel[i >= wheel_size][i % wheel_size];
}

// Use a "peg" class to represent a marker that moves along the sieve:

// Each prime has an associated "peg" that moves along the sieve.
class prime_peg
{
    Number n; // current position
    Number i; // the increment (2*p)

public:
    explicit prime_peg(Number p)
        : n{p*p},
          i{p+p}
    {}

    operator Number() const { return n; }
    prime_peg& operator++() { n += i; return *this; }

    // To use in priority_queue, sort largest first
    bool operator<(const prime_peg& other) const {
        return std::tuple(n, i) > std::tuple(other.n, other.i);
    }
};

// Set up the heap of pegs, and run the algorithm:

unsigned long find_nth_prime(unsigned long n)
{
    // no need to insert any primes higher than this
    static Number const max_peg = std::sqrt(~0UL);

    // deal with small numbers
    if (!n) throw std::invalid_argument("Requested zeroth prime");
    if (!--n) return 2;         // the only even prime
    if (!--n) return 3;         // wheel prime
    if (!--n) return 5;         // wheel prime

    // The queue m contains one peg for each prime.
    auto m = std::priority_queue<prime_peg>{};

    Number i = 7;               // first non-wheel prime
    for (auto next = prime_peg{i};  i < ~0UL/2;  i+=2) {
        if (!wheel_prime(i)) continue;
        {
            if (i < next) {
                if (!--n) break;
                if (i < max_peg) m.emplace(i);
            }
        }

        for (;  next <= i;  next = m.top()) {
            m.pop();
            do { ++next; } while (!wheel_prime(next));
            m.push(next);
        }
    }
    return i;
}

// I've made some small improvements to the test code based on other reviews:

#include <iostream>
#include <string>

int main(int argc, char **argv)
{
    std::cout.imbue(std::locale(""));
    for (int i = 1;  i < argc;  ++i) {
        try {
            std::cout << find_nth_prime(std::stoul(argv[i])) << std::endl;
        } catch (std::exception& e) {
            std::cerr << "Invalid argument: " << argv[i]
                      << ": " << e.what() << std::endl;
        }
    }
}