PROBLEM LINK:
Author: Pavel Sheftelevich
Tester: Sergey Kulik
Editorialist: Pawel Kacprzak
DIFFICULTY:
Medium
PREREQUISITES:
Merge sort, Fenwick tree
PROBLEM:
The problem can be reformulated in the following way. Given a rectangular matrix A of size N \times M, filled with integers, find out, how many submatrices of A are there, such that the sum of elements in each one of them in less or equal to given S. In every subtask, you will have to handle at most 10 test cases.
QUICK EXPLANATION:
Use fast counting method based on sorting in order to speed up the obvious O(N^2 \cdot M^2) method.
EXPLANATION:
SUBTASK 1
In the easiest subtask, N, M \leq 20, so we can iterate over all possible O(N^2 \cdot M^2) submatrices, and for each one, compute its sum in O(N \cdot M) time summing up all its elements, and check if the resulting sum is less or equal to S. This solution has time complexity O(N^3 \cdot M^3), which is enough here.
SUBTASK 2
Now, we have to deal with N, M \leq 50, so we cannot be so naive in computing the result as in the previous subtask. However, the matrix is still small enough to iterate over all possible of its submatrices. If we only are able to decide in a constant time, if a particular submatrix has a sum of its elements less or equal to S, we can solve this subtask in O(N^2 \cdot M^2) time iterating over all possible submatrices and making a fast check for each one of them.
How to decide in a constant time if a given submatrix B of A has a sum of its elements less or equal to S? Well, we can use prefix sums here.
First, let’s compute P[i][j] table, where P[i][j] is the sum of elements in a submatrix, which has its top left corner in the first row and the first column, and its bottom right corner, in the i^{th} row and in the j^{th} column. Notice that we can compute the whole P table in O(N \cdot M) time, but O(N^2 \cdot M^2) is also sufficient here. Then if we want to compute the sum of elements in a submatrix B[i_1, j_1, i_2, j_2], we can use P table to do that. Let X be the sum of elements in B, then X = P[i_2][j_2] - P[i_2][j_1 - 1] - P[i_1 - 1][j_2] + P[i_1 - 1][j_1 - 1]. Using this trick, we are able to compute the sum of any submatrix in a constant time, which leads to a O(N^2 \cdot M^2) solution.
SUBTASK 3
In the last subtask, we have N, M \leq 150, so the above solution is too slow. If we can only reduce the time complexity to something like O(N^2 \cdot M \cdot \log M), we are good to go. How to do that? Let’s first fix two rows. We will show how to count the number of submatrices with the i^{th} row as its top row and the j^{th} row as its bottom row, which have sum of its elements less or equal S. In order to get the final result, we will accumulate computed results for all pairs of top and bottom rows.
First, let’s fix some index k, which will stand for the rightmost column of a submatrices we search for. Next, rather than counting submatrices with sum at least S, we will count submatrices with sum less than S and subtract them from all possibles submatrices. In other words, we are interested in the number of submatrices ending in column k with sum less than S.
How to count this amount? Well, we can use slightly modificated P table defined in the solution to the second subtask here. Let P[k] be the sum of elements in a submatrix with k fixed as the rightmost column and the first columns as its leftmost column (of course it also has i^{th} row as its top row, and j^{th} row as its bottom row - both these rows are fixed).
Now, for a fixed k, we want to compute how many values x < k are there, such that P[k] - P[x] < S. Notice that if we have all values of P[y] for y < k sorted in ascending order, we can just find the smallest index x in this sorted order, for which P[k] - P[x] < S, and then all submatrices corresponding to indexes greater than x in the sorted order, also have the sum of elements smaller than S.
This last task can be solved beautifully with merge sort as Sergey did while he was testing the problem. Please refer to his implementation, because it explains the best how code a merge sort based solution. This solution is very similar to counting inversions in an array using merge sort, the base idea is the same here - first we divide the array of values P[y] into two halves, for which we solve the problems recursively. Then, for each index from the sorted right half, i.e. the value of P[k], we count the number of indexes x of elements in the sorted left half, for which P[k] - P[x] < S. Since both halves are sorted, we can do this in linear time. After that, we merge both halves into one sorted array as in merge sort.
Some participants uses different techniques, like for example the Fenwick Tree. Any data structure sufficient to count the submatrices as described above is perfect here.
After we solve the above subproblem for all pairs of rows, we have the result. The total complexity of this method, as intended, is O(N^2 \cdot M \cdot \log M).
AUTHOR’S AND TESTER’S SOLUTIONS:
Author’s solution can be found here.
Tester’s solution can be found here.