nth_prime: iteration 1

cpp/nth-prime/nth_prime.cpp

@@ -1,7 +1,7 @@
 #include "nth_prime.h"
 
+#include <cmath>
 #include <cstring>
-#include <iostream>
 #include <stdexcept>
 #include <vector>
 
@@ -10,29 +10,38 @@ namespace prime {
 
 using namespace std;
 
-number_t nth(size_t n)
+number_t nth(size_t n, size_t sieve_size)
 {
     if (n < 1)
-        throw domain_error("invalid prime index");
+        throw domain_error("invalid prime number index");
 
     vector<number_t> primes = {2, 3, 5, 7, 11, 13};
-    primes.reserve(n - 1);
 
-    static const size_t sieve_size = 32*1024;
+    if (n <= primes.size())
+        return primes[n - 1];
+
+    primes.reserve(n);
+
+    if (auto_sieve_size == sieve_size) {
+
+        sieve_size = min(size_t(32 * 1024),
+                         size_t(n * (log(n) + log(log(n))) + 0.5));
+    }
+
+    vector<char> sieve(sieve_size, false);
 
-    bool sieve[sieve_size] = {0};
     sieve[0] = sieve[1] = true;
 
-    size_t page = 0;
+    size_t first_number = 0;
 
-    while (primes.size() < n - 1) {
+    while (true) {
 
         for (auto p : primes) {
 
-            size_t start = ((page*sieve_size + p - 1)/p)*p
-                - page*sieve_size;
+            size_t reminder = first_number % p;
+            size_t begin = reminder ? p - reminder : 0;
 
-            for (size_t i = start; i < sieve_size; i += p) {
+            for (size_t i = begin; i < sieve_size; i += p) {
 
                 sieve[i] = true;
             }
@@ -43,20 +52,26 @@ number_t nth(size_t n)
             if (sieve[i])
                 continue;
 
-            number_t p = page*sieve_size + i;
+            number_t p = first_number + i;
             primes.push_back(p);
 
+            if (primes.size() == n)
+                return p;
+
             for (size_t j = i; j < sieve_size; j += p) {
 
                 sieve[j] = true;
             }
         }
 
-        memset(sieve, 0, sieve_size);
-        ++page;
+        size_t next_first_number = first_number + sieve_size;
+
+        if (next_first_number < first_number)
+            throw runtime_error("failed to reach prime number");
+
+        first_number = next_first_number;
+        memset(&sieve[0], false, sieve.size()*sizeof(sieve[0]));
     }
-
-    return primes[n - 1];
 }
 
 }

cpp/nth-prime/nth_prime.h

@@ -8,6 +8,11 @@ namespace prime {
 
 typedef unsigned long long number_t;
 
-number_t nth(size_t);
+enum : size_t {
+
+    auto_sieve_size = 0,
+};
+
+number_t nth(size_t n, size_t sieve_size = auto_sieve_size);
 
 }