Official

D - Pond Editorial by en_translator


We consider solving this problem with binary search. That is, we consider the following question: “is there a block of \(K\times K\) squares in the \(N\times N\) grid such that the median is less than or equal to \(X\)?” Then there exists an integer \(Y\) that for every \(X\geq Y\), the answer for the question is Yes, or if \(X<Y\), the answer is No. Now we fix some \(X\).

The median is a little bit difficult value to handle with, but by the definition this time, the question can be rephrased as “is there a block of \(K\times K\) squares in the \(N\times N\) grid, which contains \(\left\lfloor\frac{K^2}{2}\right\rfloor+1\) or more squares greater than \(X\)?” Here, “a square greater than \(X\)” means a square \((i,j)\) such that \(A_{i,j} > X\), where \((i, j)\) denotes the square in the \(i\)-th row and the \(j\)-th column.

A naive brute-force approach for this question will require \(\Theta(N^2K^2)\) time. However, all we have to do is count the number of squares satisfying the condition; therefore, if we have a grid, each square of which is assigned \(1\) if it satisfies the condition or \(0\) if not, it is equivalent to finding the sum for each block. For the problems like this, two-dimensional cumulative sums is effective. For more explanation about two-dimensional cumulative sums, you can find more detailed information on the internet too. In this problem, you can use the way described below. In fact, it can be applied for a rectangle as well as a square.

Let \(S_{i,j}\) be the number of squares \((i', j')\) in the block containing \((1,1)\) as the top-left and \((i,j)\) as the bottom-right square, such that \(A_{i',j'} >X\). How can we compute them? For convenience, we define \(S_{0,i}=S_{i,0}=S_{0,0}=0\) (which can be regarded as that of collapsed block of, for example, \(0\) rows and \(i\) columns). Then, \(S_{i, j}\) can be calculated as

\[ S_{i,j}= \begin{cases} S_{i-1,j}+S_{i,j-1}-S_{i-1,j-1}+1 & (A_{i,j}>X\ のとき) \\ S_{i-1,j}+S_{i,j-1}-S_{i-1,j-1} & (A_{i,j}\leq X\ のとき) \end{cases} \]

in the increasing order of \(i\) and \(j\). This can be easily understood by drawing diagrams.
Next, in the block having \((P,Q)\) as the top-left and \((P', Q')\) as the bottom-right square, the number of squares such that \(A_{i,j} >X\) can be found as \(S_{P',Q'}+S_{P-1,Q-1}-S_{P',Q-1}-S_{P-1,Q'}\). This can be understood by drawing a diagram too, noting in which block \(S_{i,j}\) is the number of squares that satisfy the conditions.

This enables us to find the number of squares such that \(A_{i,j} > X\) for each of \((N-K+1)\times (N-K+1)\) number of \(K\times K\) blocks. If any of them satisfies \(\left\lfloor\frac{K^2}{2}\right\rfloor+1\), then the answer is Yes; otherwise it is No. For the question where \(X\) is fixed, both finding \(S_{i,j}\) and computing each block can be processed in a total of \(O(N^2)\) time, so the problem has been solved in a total of \(O(N^2)\) time.

Obviously, the answer is between \(0\) and \(\max(A_i)\); let \(ng=-1,\ ok=\max(A_i)\), and repeat the following procedure until \(ng+1=ok\): let \(X=(ng+ok)/2\), and update \(ok=X\) if the answer is Yes, or if No, \(ng=X\). Then the answer is the final \(ok\). Since this operation finishes in at most \(\lfloor\log_2 (\max(A_i))\rfloor+1\) repetitions, the original problem has been solved in a total of \(O(N^2\log (\max(A_i)))\) time. This is fast enough.

Sample code in C++:

#include <bits/stdc++.h>
using namespace std;
 
#define N 805
#define MAX_A 1000000000
#define rep(i, n) for(int i = 0; i < n; ++i)
 
int main(void) {
	int n, k, lim;
	int a[N][N];
	int s[N][N];
	rep(i, N) {
		s[i][0] = 0;
		s[0][i] = 0;
	}
	int ng = -1;
	int ok = MAX_A;
	int mid;
	bool ext;
	cin >> n >> k;
	lim = ((k*k) / 2) + 1;
	rep(i, n) {
		rep(j, n)cin >> a[i][j];
	}
	while ((ng + 1) < ok) {
		mid = (ng + ok) / 2;
		rep(i, n) {
			rep(j, n) {
				s[i + 1][j + 1] = s[i + 1][j] + s[i][j + 1] - s[i][j];
				if (a[i][j] > mid)s[i + 1][j + 1]++;
			}
		}
		ext = false;
		rep(i, n - k + 1) {
			rep(j, n - k + 1) {
				if ((s[i + k][j + k] + s[i][j] - s[i][j + k] - s[i + k][j]) < lim)ext = true;
			}
		}
		if (ext)ok = mid;
		else ng = mid;
	}
	cout << ok << endl;
	return 0;
}

posted:
last update: