前缀和，二分 - [ABC203D] Pond

AT_abc203_D Pond

题意

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

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

思路

1. 暴力枚举

暴力枚举右下端点，将 \(k\times k\) 矩阵中的元素取出，然后排序。时间复杂度为 \(O(n^2k^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\) 的个数小于等于 $\left \lfloor \frac{k^2}{2} \right \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;
}