CLSUMG - Editorial



Author: Avijit Agarwal
Tester and Editorialist: Soumik Sarkar




Prime sieve, Fast fourier transform


Let f(x)= the number of pairs (a, b) such that 0 \le a \le b < x and a and b are prime and x = a + b. Given a number N, find the number of pairs (a, b) such that 0 \le a \le b < x and f(N) = f(a) + f(b).


This problem was inspired by Goldbach’s conjecture :smiley:

Let us assume can find f(x) in constant time for any given x. Then for a query N naively trying all pairs will result in a complexity of \mathcal{O}(N^2). Instead we can build up a frequency map/array of the sequence f by upto index N-1. Let’s call this frequency map cnt. Now we can count how many pairs add up to f(N) by iterating over suitable i and adding cnt(i) \times cnt(f(N)-i) to the total. Note that we are looking for pairs where a \le b so we must make sure not to double count any pair. The complexity is \mathcal{O}(N) for one query.

This approach relies on the availability of all f(x) in constant time. One way to achieve that would be to calculate all primes upto the required range beforehand using a sieve. Let there be P such primes, then we can compute the array f in \mathcal{O}(P^2) time, but that would be too slow.

For an efficient solution, consider the polynomial P(x) = x^2 + x^3 + x^5 + .... + x^p where in each term the exponent of x is a prime and p is the last prime we need to be concerned with.

P(x)^2 = ( x^2 + x^3 + x^5 + ....)^2 = x^4 + 2x^5 + x^6 + 2x^7 + 2x^8 + ...

Multipying x^a and x^b gives x^{a+b}. Thus in the polynomial above the coefficient of x^k is the number of ways to represent k as a sum of 2 primes. Note that here the pairs (a, b) and (b, a) for a \ne b are counted separately, so the values must be adjusted to suit the given definition of f(x).

Squaring a polynomial of degree n naively takes \mathcal{O}(n^2) time. However there exists the awesome method of Fast Fourier Transform which can be used to efficiently compute the product of two polynomials in \mathcal{O}(n \log n).

The complexity is \mathcal{O}(N \log N + TN).


Author’s solution can be found here
Tester’s solution can be found here.


Never knew we can use fft so easily in python! Can you explain a bit or give some link for understanding the functions involved?

You can refer to the documentation for numpy.fft and scipy.fftpack. There is one useful function in scipy.signal, which is fftconvolve that I have used.

1 Like