前缀和，二分 - [ABC203D] Pond

AT_abc203_D Pond

题意

给出一个 $N\times N$ 的矩阵，然后依次枚举 $k\times k$ 的子矩阵。

对于 $k\times k$ 的子矩阵，一共有个 $k^2$ 元素，找出其中的中位数。这里的中位数是子矩阵中元素从大到小排列的第 $\lfloor \frac{k^2}{2}\rfloor$ 个元素。问对于所有子矩阵的中位数中，最小的那个中位数是多少。

思路

1. 暴力枚举

暴力枚举右下端点，将 $k\times k$ 矩阵中的元素取出，然后排序。时间复杂度为 $O(n^2{k}^2\log k)$，绝对超时。

2. 二分答案+二维前缀和

我们不妨转变思路，从枚举矩阵到二分中位数。判断一个值能否成为一个 $k\times k$ 矩阵的中位数。

如何判断是否为中位数？

假设一个数 $x$ ，如果 $a_{i,j}>x$，那么另一个二维矩阵中的 $b_{i,j}$ 便设为 $1$，否则设为 $0$。这样便将矩阵 $a$ 转变成了一个零一矩阵 $b$。对其求前缀和，这样就可以求得 $b$ 中任意一个 $k\times k$ 矩阵中 $1$ 的个数。如果其中 $1$ 的个数小于等于 $\lfloor \frac{k^2}{2}\rfloor$ 那么将上界调为 $x$，否则将下界调整为 $x+1$。时间复杂度为 $O(n^2\log n)$。

Code

#include<bits/stdc++.h>
using namespace std;
const int MAXN = 800 + 5;
int n, k, l, r, kpow, lim;
int a[MAXN][MAXN], b[MAXN][MAXN];
bool check(int x) {
	int z;
	for (int i = 1; i <= n; i++) {
		for (int j = 1; j <= n; j++) {
			b[i][j] = a[i][j] > x;//转化为01矩阵
			b[i][j] = b[i][j - 1] + b[i - 1][j] - b[i - 1][j - 1] + b[i][j];//求二维前缀和
		}
	}
	bool flag = 0;
	for (int i = 1; i <= n - k + 1; i++) {
		for (int j = 1; j <= n - k + 1; j++) {
			z = b[i + k - 1][j + k - 1] - b[i - 1][j + k - 1] - b[i + k - 1][j - 1] + b[i - 1][j - 1];//一的个数
			if (z <= kpow) {
				flag = 1;//如果找到了一个可能的值就退出
				break;
			}
		}
		if (flag) break;
	}
	return flag;
}
signed main() {
	scanf("%d%d", &n, &k);
	kpow = k * k / 2;
	for (int i = 1; i <= n; i++) {
		for (int j = 1; j <= n; j++) {
			scanf("%d", &a[i][j]);
			r = max(a[i][j], r);
		}
	}
	while (l < r) {//二分
		int mid = (l + r) >> 1;
		if (check(mid)) {//可行，找更小的
			r = mid;
		} else {
			l = mid + 1;
		}
	}
	printf("%d", r);
	return 0;
}