improved prime number generator to check only up to sqrt(n) instead of n (#1984)

* improved prime number generator to check only up to sqrt(n) instead of n

* added old version as slow_primes() and named new, faster version primes()

* fixed docstring in slow_primes

* Add a timeit benchmark

* Update prime_numbers.py

Co-authored-by: Christian Clauss <[email protected]>

Author: steven7214

maths/prime_numbers.py

@@ -1,4 +1,32 @@
 from typing import Generator
+import math
+
+
+def slow_primes(max: int) -> Generator[int, None, None]:
+    """
+    Return a list of all primes numbers up to max.
+    >>> list(slow_primes(0))
+    []
+    >>> list(slow_primes(-1))
+    []
+    >>> list(slow_primes(-10))
+    []
+    >>> list(slow_primes(25))
+    [2, 3, 5, 7, 11, 13, 17, 19, 23]
+    >>> list(slow_primes(11))
+    [2, 3, 5, 7, 11]
+    >>> list(slow_primes(33))
+    [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31]
+    >>> list(slow_primes(10000))[-1]
+    9973
+    """
+    numbers: Generator = (i for i in range(1, (max + 1)))
+    for i in (n for n in numbers if n > 1):
+        for j in range(2, i):
+            if (i % j) == 0:
+                break
+        else:
+            yield i
 
 
 def primes(max: int) -> Generator[int, None, None]:
@@ -21,7 +49,9 @@ def primes(max: int) -> Generator[int, None, None]:
     """
     numbers: Generator = (i for i in range(1, (max + 1)))
     for i in (n for n in numbers if n > 1):
-        for j in range(2, i):
+        # only need to check for factors up to sqrt(i)
+        bound = int(math.sqrt(i)) + 1
+        for j in range(2, bound):
             if (i % j) == 0:
                 break
         else:
@@ -32,3 +62,8 @@ def primes(max: int) -> Generator[int, None, None]:
     number = int(input("Calculate primes up to:\n>> ").strip())
     for ret in primes(number):
         print(ret)
+
+    # Let's benchmark them side-by-side...
+    from timeit import timeit
+    print(timeit("slow_primes(1_000_000)", setup="from __main__ import slow_primes"))
+    print(timeit("primes(1_000_000)", setup="from __main__ import primes"))