PROBLEM LINK
DIFFICULTY
HARD
PREREQUISITES
Dynamic Programming, Discrete Fourier Transform, Fast Fourier Transform
PROBLEM
You are given a list of integers between 1 and 30000, inclusive. You have to find triplets i < j < k in A[], such that A[i], A[j] and A[k] are in Arithmetic Progression.
QUICK EXPLANATION
The trivial algorithm to solve this problem is O(N3). This is of course too slow for N = 100,000.
Let us consider the O(N2) algorithm below to solve this problem. It will also be slow, but it gives us ideas that can lead to a solution.
prev[1 ... 30000] = {0} // we will maintain frequencies here
result = 0
for j = 1 to N // fix j
for k = j+1 to N // fix k
Ai = 2*A[j] - A[k] // the value at i
result += prev[Ai] // add the number of choices for i
prev[ A[j] ]++
Of course this is too slow.
EXPLANATION
Let us consider an even slower algorithm. The reason to consider this algorithm, is because it introduces a very powerful idea to solve the problem. Once we understand how that works, we will talk about how we can solve the problem.
prev[1 ... 30000] = {0} // we will maintain frequencies here
next[1 ... 30000] = {0} // we will mataiain frequencies here too! This time,
// we maintain frequencies of occurance 'after' j
for v = 1 to N
next[ A[v] ]++
for j = 1 to N
next[ A[j] ]--
cnt[1 ... 60000] = {0}
for Ai = 1 to 30000
for Ak = 1 to 30000
cnt[Ai + Ak] += prev[Ai] * next[Ak]
result += cnt[ 2*A[j] ]
prev[ A[j] ]++
Notice how we maintain cnt. This is actually equivalent to multiplication of two polynomials!
Suppose we had two polynomials
- P = p1x + p2x2 + p3x3 … + p30000x30000
- Q = q1x + q2x2 + q3x3 … + q30000x30000
The product of these two polynomials R, can be calculated as follows
for i = 1 to 30000
for j = 1 to 30000
R[i+j] = P[i] * Q[j]
We can optimize the product of the two polynomials by using Discrete Fourier Transforms!
fP = DFT(P)
fQ = DFT(Q)
fR = {0}
for i = 1 to 216
fR[i] = fP[i] * fQ[i]
R = iDFT(fR)
Where, fR is the Discrete Fourier Transform of R, and iDFT is the inverse Discrete Fourier Transform method.
Discrete Fourier Transforms, and inverse Discrete Fourier Transforms of an array of size N can be calculated in O(N log N) time using the Fast Fourier Transform method! The link to the same has been provided in the PREREQUISITES section above.
Hence we can modify our O(N*300002) algorithm into a O(N*30000*16) algorithm.
This, is still too slow to solve the problem. Now we introduce the idea that solves the problem.
Optimization using Blocks
Suppose we split the entire list into K blocks each of size N/K.
- prev[] represents the frequency of the numbers in blocks before the current block.
- next[] represents the frequency of the numbers in blocks after the current block.
- inside[] represents the frequency of numbers within the current block.
We have three cases to consider while choosing triplets.
- The entire triplet lies in the current block
- This can be calculated in O(N2/K2) time similar to the first O(N2) algorithm we described. We will be using the inside array for the frequencies.
-
Two numbers in the triplet lie in the current block
- This can also be calculated in O(N2/K2) time.
- We use the first O(N2) algorithm. We just use prev / next instead of inside.
- Only one number lies in the current block
- We find the product of the two polynomials prev and next using FFT.
- This step takes O(30000*16) calculations.
This the overall complexity of the algorithm would be
O(K*30000*16 + N*N/K)
You can experiment and choose an optimal K. The tester’s solution chooses K = 30, the setter’s solution chooses K = 40.
This will fit within the time limit.
SETTERS SOLUTION
Can be found here.
TESTERS SOLUTION
Can be found here.
