def combine_counts(a, b, mod):
n = a.size
if n == 2:
return np.array([a[0]*b[0]+a[1]*b[1], a[0]*b[1]+b[0]*a[1]]) % mod
elif n == 1:
return (a*b) % mod
a %= mod
b %= mod
pivot = n/2
u = combine_counts(a[:pivot]+a[pivot:], b[:pivot]+b[pivot:], mod)
v = combine_counts(a[:pivot]-a[pivot:], b[:pivot]-b[pivot:], mod)
res = np.zeros(n, dtype = a.dtype)
res[:pivot] = ((u + v) * pe.mod_inv(2, mod)) % mod
res[pivot:] = ((u - v) * pe.mod_inv(2, mod)) % mod
return res
def solve(n):
primes = np.array(pe.primes_and_mask(n)[0], dtype = np.int64)
pi = len(primes)
#product of first k primes mod p, for each prime
mod_prods = 2*np.ones(pi, dtype=np.int64)
a_mods = np.ones(pi, dtype = np.int64)
#a_0 is 0, I guess
a = 1
mask = primes == -1
for i in xrange(1, pi):
#we want to find the smallest k with k*p_1*...*p_i + A_i = t*p_(i+1) + (i+1)
#solving this equation mod p_(i+1) for k yields the answer
p = primes[i]
k = ((i+1 - a_mods[i]) * pe.mod_inv(mod_prods[i], p)) % p
a_mods[i:] = (a_mods[i:] + k*mod_prods[i:]) % primes[i:]
mod_prods[i:] *= p
mod_prods[i:] %= primes[i:]
mask[a_mods==0] = True
print np.sum(primes[mask])
