PROBLEM LINK:
Author and Editorialist: Hussain Kara Fallah
Tester: Michael Nematollahi
PREREQUISITES:
divide and conquer, dynamic programming
PROBLEM:
Let’s call an integer sequence nice if the sum of its elements is divisible by a given integer K. You are given an integer sequence A_1, A_2, \ldots, A_N.
For an arbitrary non-empty sequence of indices i_1 \lt i_2 \lt \ldots \lt i_M, let’s call a subsequence A_{i_1}, A_{i_2}, \ldots, A_{i_M} very nice if it is nice and i_M-i_1 \ge W. Find the number of very nice subsequences modulo 10^9+7.
EXPLANATION:
If approaching the array as a complete piece is kind of hard. Try to use divide and conquer. In our problem, it works pretty fine. (There were some solutions without divide and conquer but probably harder to come up with).
Let’s suppose we are processing a part of the array A[L\,...\,R].
Break it into 2 pieces, A[L\,...\,mid] and A[mid+1\,...R]. where (mid = \frac{L+R}{2})
Subsequences completely residing in the left part can be calculated by passing the left part to our divide and conquer function. We should do the same for the right part also.
Now we should calculate the very nice subsequences that have at least one element in the left part and one element in the right part.
Let’s introduce 2 functions:
F_{left}(pos, m) denotes the number of subsequences considering numbers in the subarray A[pos,mid] and have their sum modulo K equal to exactly m.
F_{right}(pos, m) denotes the number of subsequences considering numbers in the subarray A[mid+1,pos] and have their sum modulo K equal to exactly m.
calculating both of tables is a simple DP exercise that’s left to you. Calculating both tables take O((R-L+1)*K) in total.
Now let’s find our nice subsequences.
Fix the rightmost element (which has index i_M in problem description). For each x in range [0,K-1] let’s find the number of subsequences in the right part such that the last element is our fixed element and their sum modulo K is equal to x. Their number is F_{right}(i_M,x)-F_{right}(i_M-1,x)
(consider processing one possible value of x as the process will be the same for the rest).
Now we want to combine these subsequences with some subsequences from the left part (they must have their sum modulo K equal to y=K-x). The first element must be at least W elements away from the last one. Since we fixed the last one, our life is much easier now. The first element has index i_1 <= i_M-W.
By the same logic, what is the number of subsequences in the left part such that the first element has the index i_1 which is at most i_M-W and have sum modulo K equal to y. Well it’s equal to F_{left}(L,y)-F_{left}(i_M-W+1,y)
So for a fixed last element i_M and a fixed modulo sum x if the right part’s subsequence
Result \, += \, (F_{right}(i_M,x)-F_{right}(i_M-1,x)) \, * \, (F_{left}(L,y)-F_{left}(i_M-W+1,y))
Complexity: O(N*K*log\,N)