PROBLEM LINK:
Author: Rahim Mammadi
Primary Tester: Amir Reza PoorAkhavan
Editorialist: Hussain Kara Fallah
DIFFICULTY:
Easy-Medium
PREREQUISITES:
Dynamic Programming
PROBLEM:
You are given an integer sequence A_1,A_2,…,A_N and integers K and M. For 1≤i≤j≤N, let’s define S(i,j) as the number of ways to choose exactly K elements of the contiguous subsequence A_i,A_{i+1},…,A_j in such a way that the median of these K elements is ≥M.
Find the sum of S(i,j) over all i,j such that 1≤i≤j≤N. Since this sum may be large, calculate it modulo 10^9+7.
Constraints:
1 \leq N \leq 10^5
1 \leq K \leq 100
K is odd
EXPLANATION:
Let’s replace every element which is less than M by 0 and the rest by 1. Now we are currently interested in subsequences where the number of ones is strictly more than half of K.
Let’s first solve an easy problem and count only the number of subsequences such that their median is \geq M. Our task turns out to count the number of subsequences consisting of K elements and their sum is strictly larger than \frac{K}{2}.
We can use a naive dynamic programming approach. Let’s denote F(Pos,T,S) as the number of subsequences ending up to the element at index Pos and consisting of T elements and the total subsequence sum is exactly S.
Our recurrence will look like:
F(pos,T,S) \, = \, F(pos-1,T,S) + F(pos-1\,,T-1\,,S-A_{pos})
The first trick we should do here is dropping one dimension out of the DP table because the straight-forward storage method consumes too much memory. You can notice that for a fixed pos (position), the values of the table only depends on the values of the exact previous position.
The second thing that we should notice is that once the sum is equal to or exceeds 50 then the median is guaranteed to be \geq M. So we can even optimize by assuming that the maximum possible value of S in the space is 50. We only need then to add one extra transition in one edge case which is appending an element with value 1 to a subsequence with sum already bigger or equal to 50.
Theoretically, This runs in O(N*K*K/2) which is around 5*10^8 operations. But it’s far less than that by doing simple breaks and avoiding unnecessary computations (states in space where the answer of them is 0). A solution with such complexity runs in ~1 second.
Till now, we haven’t solved the original problem yet. Notice that a subsequence starting at position L and ending at position R should be counted exactly L \, * \, (N-R+1) times in the answer. (Obviously, any contiguous interval beginning at position L or before and ending at R or after contains this subsequence).
By setting F(pos,1,A_{pos})\,=\,pos , we can guarantee that each subsequence starting at any position L will be counted L times. Now we can easily obtain the number of subsequences with any fixed (T,S) ending at any position in the array, so for each position R we check the number of valid subsequences ending there and multiply the number by (N-R+1) and add everything to our final answer. And don’t forget the modulo.