Tags: prefix sums
The first thing you should think when you see this problem is prefix sums because it looks like this problem can be done in O(n*n).
1. Count the number of white lines there already are (like in WWW there's already 1 white line)
2. Create 2 arrays, a 2D horizontal prefix sum array and a 2D vertical prefix sum array.
3. For each horizontal line from the left then, if it's at most K places away from becoming a line horizontally, mark that position as 1.
Example 1:
BBB
WBB
BBB
If K = 2, the horizontal sum array (before you spread the sums to be a prefix sum array) would be
0 0 0
0 1 0
0 0 0
If K = 3, the horizontal sum array would be
1 1 1
1 0 0
1 0 0
4. For each vertical line from the top, if it's at most K places away from becoming a line vertically, mark that position as 1.
Example 2:
If K = 1, the vertical sum array would be
0 0 0
0 1 0
0 0 0
If K = 2, the vertical sum array would be
0 1 0
0 1 0
0 0 0
5. Then, you spread the sums. For the horizontal sum array, you spread the sums vertically and for the vertical you spread them horizontally.
Like the K = 2 array from Example 1, when you spread it vertically, it looks like this:
0 0 0
0 1 0
0 1 0
And for the K = 1 array from Example 2, you spread it like this:
0 0 0
0 1 1
0 0 0
You find out why we do it this way in step 6.
6. Iterate through all the cells where you can put the eraser (
i = 1, 2, 3... N-K,
j = i, 2, 3... N-K). From the horizontal sum array, get the vertical sum of the column j. Then from the vertical sum array, get the horizontal sum of the row i. These sums add to the number of lines because they are the number of lines that need at most K more cells to form lines. I think it's easier to understand in an example.
From example 2, we have
BWB
BBB
BWB
Let's say K = 2.
The vertical sum array will be
0 1 0
0 1 0
0 0 0
The horizontal sum array will be
0 0 0
0 0 0
0 0 0
Then, after we spread the sums, the vertical sum array will be
0 1 1
0 1 1
0 0 0
Then, we check every cell. If we check cell (0, 1), the sum we get is 1.
0
1 0 (getting the sum of 0+1 from the prefix array above)
0 1 0
0 0 0
Which is 1, and we can get this from the prefix sum array (the one above this array) by doing arr[0][2]-arr[0][0].
Another example:
BWB
WBW
BWB
K = 1
The vertical array will be
0 0 0
0 1 0
0 0 0
The horizontal array will be
0 0 0
0 1 0
0 0 0
Spreading the sums, the vertical array will be
0 0 0
0
1 1
0 0 0
The horizontal array will be
0 0 0
0
1 0
0 1 0
When we try (1, 1), the sum we get is in the bolded parts above (we only go 1 far since K = 1).
My submission:
https://codeforces.com/contest/1200/submission/58621782