OPDIVF - Editorial


Contest - Division 3
Contest - Division 2
Contest - Division 1

Setter: Jatin Khandual
Tester: Anshu Garg




FFT/NTT Divide and Conquer


Given the prime factorisation of a positive integer N, determine the sum of all positive integers D such that

  1. N is divisible by D and,
  2. the exponents of the prime factors of D add up to K.


Let A_i = \{p_i^0, p_i^1,\dots,p_i^{q_i}\} for all 1\le i\le m.
Then, we are interested in finding the sum of all D where

  • D is the product of exactly one element from each A_i and,
  • the sum of the exponents of the selected elements equals K.
Naive DP (TLE; Differs from the actual solution)

Let dp[j][k] be the sum of all D such that

  • D is the product of exactly one element from each A_i\space(i\le j)
  • the sum of the exponents of the selected elements equals k.

This is a classical DP problem, whose recurrence relation is:

dp[j][k] = \sum_{x=0}^{\min(k,q_j)} dp[j-1][k-x]*p_j^x

with the (obvious) base case dp[0][0] = 1.

The time complexity to compute the above recurrence using DP is O(mK^2), too much for this problem, to run within the time limit.

The key to solving the problem is modelling it as an FFT one.
Define polynomial equations F_i(x) for all 1\le i\le m as:

F_i(x) = p_i^0 + p_i^1x + p_i^2x^2+\dots+p_i^{q_i}x^{q_i} = \sum_{j=0}^{q_i} p_i^jx^j

It is not hard to see that for a particular K, the coefficient of x^K in the polynomial product F_1(x)*F_2(x)*\dots*F_m(x) is the desired answer.

Thus, we have reduced the problem to multiplying m polynomial equations, which can be done efficiently using FFT/NTT (Since FFT has precision errors over multiplication of large numbers; NTT should be used here).

Multiplication of two polynomials using FFT takes O(S\log S), where S is the maximum degree of the two.
Multiplying m polynomials iteratively therefore, has worst case O(m*T\log T) time complexity, where T is the sum of the degrees of all F_i. This naive approach will TLE.

Instead, we use divide and conquer to calculate the product of the polynomials using FFT/NTT.


Let S be the set of polynomials we want to multiply. Also, let T be the sum of the degrees of all polynomials in S.

The divide and conquer algorithm looks as follows:

	n = len(S)
	if n == 1:
	    return S[0]
	return multiply(S[:n/2]) * multiply(S[n/2:])

In the above pseudo code, we split S into two sets, recursively calculate their products, and then multiply the two products and return the value.

Since the sum of the degrees of the multiplied polynomials at each level in the recursion is T, and the number of levels is \log |S|, the time complexity of this approach is:

O(T\log T*\log |S|)

which is sufficient to AC the problem, for the given constraints.


Using divide and conquer to multiply m polynomials using NTT takes

O(T\log T\log m)

where T is equal to the sum of all q_i.


Editorialist’s solution can be found here

Experimental: For evaluation purposes, please rate the editorial (1 being poor and 5 excellent)

  • 1
  • 2
  • 3
  • 4
  • 5

0 voters