Fibonacci Primitve Roots of Primes (Project Euler)

I had heard about Project Euler for a few years now but it was not until just a few days ago that I decided to give it a shot. And I quickly realized that solving problems on it is a lot of fun. On my first go at it, I solved eleven problems in a row. Once I got to problem 18 I realized that perhaps it would make more sense for me to start from the end instead of the beginning, given that the problems at the beginning are rather easy. So I jumped to the last one, which was, at the time of writing this, problem 437. And I was pleasantly surprised. This problem was just the right amount of difficult--difficult enough not to be trivial but not so difficult that it would take hours and hours to solve! Of course my rule for solving Project Euler problems is that I can only use Python and basic Python packages. No mathematics packages or software, no use of algorithms I did not fully understand, etc.

First, here's the problem statement, rephrased in a more mathematical language:

$$x$$ is said to be a Fibonacci primitive root for prime $$p$$ iff it is a primitive root for $$p$$ and $$x^n+x^{n+1} = x^{n+2}$$ for all integers $$n.$$ Find the sum of all primes less than $$10^8$$ with at least one Fibonacci primitive root.

My first thought on this was to first list all primes less than $$10^8$$, then for each of them list all the primitive roots, and then check each primitive root to see if it is a Fibonacci one or not. Of course, this would be excruciatingly slow for that many primes. So I decided to do some math before jumping to programming first.

The Fibonacci requirement on the primitive root can be simplified to $$x^2 - x - 1 = 0.$$ This equation has solutions $$\frac{1}{2}(1 \pm \sqrt{5}).$$ Hence for a prime to have a Fibonacci primitive root it must have $$5$$ as a quadratic residue. By the quadratic reciprocity theorem, we only need to consider primes of the form $$5k \pm 1.$$ This narrows down our search. Of course $$5$$ is a special case and I just checked it manually. Since it has a Fibonacci primitive root ($$3$$) it should be included in the answer too.

The next step would be to check if either of the two solutions above are primitive roots. For that of course, we first need to calculate the solutions. This requires calculating the square root of $$5$$ modulo the prime $$p.$$ If $$p$$ is $$3$$ modulo $$4$$ then this is easy as for such $$p$$ we have

\begin{equation*} \sqrt{5} = \pm 5^{\frac{p+1}{4}}. \end{equation*}

But for $$p=4k+1$$ we need something more sophisticated. There are two relatively efficient algorithms for this: Tonelli-Shanks and Cipolla's algorithm. I chose Tonelli-Shanks because the running time on average for this particular application is better.

Of course we still have to list all primes that are $$\pm 1$$ modulo $$5.$$ For this I used Rabin-Miller primality testing since I had already written the code for it in my post on RSA. Of course a pre-calculated list of primes could also be used for a faster algorithm.

Here is the final code:

def square_root(n, p):
"""Finds the square root of n modulo prime p. Uses Tonelli-Shanks algorithm."""
# First the easy, obvious case:
if p % 4 == 3:
return pow(n, (p + 1) / 4, p)

q, s = (p - 1) / 4, 2
while q % 2 == 0:
q /= 2
s += 1

# Find a quadratic non-residue z
z = 0
for z in prime_generator(p - 1):
break
if z == 0:
# This should never happen if p is prime and n is a quadratic residue mod p.
raise ValueError()
c = pow(z, q, p)
r = pow(n, (q + 1) / 2, p)
t = pow(n, q, p)
m = s
while t != 1:
i = 1
t1 = (t * t) % p
while t1 != 1:
t1 = (t1 ** 2) % p
i += 1
b = pow(c, 2 ** (m - i -1), p)
r = (r * b) % p
t = (t * b ** 2) % p
c = (b ** 2) % p
m = i

return r

def fibonacci_primitive_root(p):
"""
Returns a Fibonacci primitive root if p has one and None otherwise.
Assumes p is a prime and that p is congruent to 1 or -1 modulo 5.
"""
# The Fibonacci root is a root of the polynomial x^2-x-1.
# The roots of this polynomial are given by 1+sqrt(5) and 1-sqrt(5)
# So the basic idea is to simply find the square root, and check to see if either
# of those values are a primitive root.

q = square_root(5, p)
half = (p + 1) / 2

q1 = half + half * q
q2 = half - half * q

q1 %= p
q2 %= p

if is_primitive_root(q1, p):
return q1
elif is_primitive_root(n - q1, p):
return q2
return None

def sum_primes_with_fib_primitive_root(N):
s = 5L  # manually checked five. It has a Fibonacci primitive root of 3.
l = []
for n in xrange(10, N, 5):
for i in [-1, 1]:
if rabin_miller_is_prime(n + i, k=10):
root = fibonacci_primitive_root(n + i)
if root:
s += n + i
print s