Finding consecutive primes matching p, p + 4, p + 6, p + 10, p + 12, p + 16

Question

Someone asked about this question on the main StackOverflow. The full question is:

Given a value N, find p such that all of [p, p + 4, p + 6, p + 10, p + 12, p + 16] are prime.

• The sum of [p, p + 4, p + 6, p + 10, p + 12, p + 16] should be at least N.

My thinking is:

• Sieve all primes under N
• Ignore primes below (N-48)/6
• Create consecutive slices of length 6 for the remaining primes.
• Check if the slice matches the pattern.

Here's my solution. I'd appreciate some feedback.

from itertools import dropwhile, islice
def get_solutions(n):
    grid = [None for _ in range(n+1)]
    i = 2
    while i < n+1:
        if grid[i] is None:
            grid[i] = True
            for p in range(2*i, n+1, i):
                grid[p] = False
        else:
            i += 1
    sieve = (index for index, b in enumerate(grid) if b)
    min_value = (n - 48) / 6
    reduced_sieve = dropwhile(lambda v: v < min_value, sieve)
    reference_slice = list(islice(reduced_sieve, 6))
    while True:
        try:
            ref = reference_slice[0]
            differences = [v - ref for v in reference_slice[1:]]
            if differences == [4, 6, 10, 12, 16]:
                yield reference_slice
            reference_slice = reference_slice[1:] + [next(reduced_sieve)]
        except StopIteration:
            break


n = 2000000

print(next(get_solutions(n)))

# or for all solutions
for solution in get_solutions(n):
    print(solution)

Answer 1

Overall, good first question! The first obvious improvement is to factor out the code that generates the primes

def prime_sieve(n):
    grid = [None for _ in range(n+1)]
    i = 2
    while i < n+1:
        if grid[i] is None:
            grid[i] = True
            for p in range(i*i, n+1, i):
                grid[p] = False
        else:
            i += 1
    return (index for index, b in enumerate(grid) if b)

Note that for p in range(i*i, n+1, i): starts later than the for p in range(2*i, n+1, i): which you used. This is safe because anything less than the current prime squared will have already been crossed out. This difference alone makes the code about 2x faster for n = 4000000.

By separating the sieve, it makes things like profiling much easier, and you can see that most of the time this method takes is still in the sieve. Using some tricks from Find primes using Sieve of Eratosthenes with Python, we can focus our efforts on speeding this part up.

def prime_sieve(n):
    is_prime = [False] * 2 + [True] * (n - 1) 
    for i in range(int(n**0.5 + 1.5)): # stop at ``sqrt(limit)``
        if is_prime[i]:
            is_prime[i*i::i] = [False] * ((n - i*i)//i + 1)
    return (i for i, prime in enumerate(is_prime) if prime)

This prime sieve works pretty similarly, but is shorter, and about 4x faster. If that isn't enough, numpy and https://stackoverflow.com/questions/2068372/fastest-way-to-list-all-primes-below-n/3035188#3035188 can come to the rescue with this beauty which is another 12x faster.

import numpy
def prime_sieve(n):
    """ Input n>=6, Returns a array of primes, 2 <= p < n """
    sieve = numpy.ones(n//3 + (n%6==2), dtype=numpy.bool)
    for i in range(1,int(n**0.5)//3+1):
        if sieve[i]:
            k=3*i+1|1
            sieve[       k*k//3     ::2*k] = False
            sieve[k*(k-2*(i&1)+4)//3::2*k] = False
    return numpy.r_[2,3,((3*numpy.nonzero(sieve)[0][1:]+1)|1)]

At this point, further speedup would need to come from fancy number theory, but I'll leave that for someone else.

Answer 2

I agree with the tips given in the answer by @OscarSmith regarding the prime sieve. In addition, here are a few more comments on how to make the rest of your code slightly better.

• You can go one step further and create a primes_in_range function that does the dropping for you. This might actually become a bottleneck at some point, (if a is large and a >> b - a), but at that point you only need to change the implementation of one function.

• Your reference_slice is always the same length, but if you want to add the next element you need to resort to list slicing and list addition, both of which might create a copy. Instead, you can use collections.deque with the optional argument maxlen. If you append another element, the first one will be automatically pushed out.

• Similarly, you can avoid the slice in the calculation of the differences by just including the first element. You just have to add a 0 as first element in the correct differences list. Alternatively you could just define a diff function and adjust the correct differences to be relative to the previous elements instead of the first.

• You should limit your try block as much as possible. This increases readability on one hand, and reduces unwanted exceptions being caught on the other (not such a big risk here, but true in general).

• But even better is to not need it at all. A for loop consumes all elements of an iterable, which is exactly what we need here. This way we don't even need to special case the first five elements, because if the list is too short, the pattern cannot match. Alternatively, you could implement a windowed(it, n) function, which is what I have done below.

• You should try to avoid magical values. In your code there are two. The first is the length of the pattern and the second is the pattern itself. Fortunately, the former is already given by the latter. I would make the pattern an argument of the function, which makes it also more general. At this point a less generic name might also be needed (although the name I propose below is maybe not ideal either).

• You should protect the code executing your functions with a if __name__ == "__main__"; guard to allow importing from this script without it being run.

All of these changes make your function doing the actual work a lot shorter and IMO more readable. The functions I defined below are also nice things to have in you toolkit, since they are generally applicable and not only for this specific usecase.

from collections import deque
from itertools import dropwhile, islice


def prime_sieve(n):
    ...

def primes_in_range(a, b):
    yield from dropwhile(lambda x: x < a, prime_sieve(b))

def diff(it):
    it = iter(it)
    previous = next(it)
    for x in it:
        yield x - previous
        previous = x

def windowed(it, n):
    d = deque(islice(it, n), maxlen=n)
    assert len(d) == n
    for x in it:
        yield tuple(d)
        d.append(x)

def find_primes_matching_pattern_below(diff_pattern, n):
    min_value = (n - 48) / 6   # A comment explaining this bound
    primes = primes_in_range(min_value, n)
    for candidate in windowed(primes, len(diff_pattern) + 1):
        if tuple(diff(candidate)) == diff_pattern:
            yield candidate


if __name__ == "__main__":
    n = 2000000
    # from (4, 6, 10, 12, 16) relative to the first element
    diff_pattern = (4, 2, 4, 2, 4)
    print(next(find_primes_matching_pattern_below(diff_pattern, n)))

For the given testcase, your code takes 741 ms ± 34.3 ms, while this code takes 146 ms ± 3.77 ms (using the numpy prime sieve from the answer by @OscarSmith) on my machine.