[NOIP 2024 模拟2]矩阵学说

[NOIP 2024 模拟2]矩阵学说

题意

给出 $n$ 行 $m$ 列的矩阵，第 $i$ 行第 $j$ 列的元素为 $a_{i,j}$，找出满足以下条件的三元组 $(i,j,x)$ 的 数量：

1. $1\le i\le n$, $1\le j\le m$, $1\le x\le min(n-i+1,m-j+1)$
2. 矩阵的左上角 $(i,j)$ 到右下角 $(i+x-1,j+x-1)$ 恰好含 $k$ 个不同的整数。

$1\le a_{i,j}\le 100$ $1\le n,m\le 2000$

思路

考虑求出包含数的种类 $\le k$ 的个数和 $<k$ 的个数，相减得到答案。

枚举左上角，矩阵包含数的种类个数显然单调不递减，二分边长即可。

快速求区间数的种类可以用二维 ST 表加 bitset。

时间复杂度：$O(nm\log min(n,m))$。

代码

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int N = 1505;
int n, m, k, a[N][N], lg[N];
bitset <100> st[N][N][12];
void init() {
	lg[1] = 0;
	for (int i = 2; i <= n; i ++) lg[i] = lg[i >> 1] + 1; 
	for (int i = 1; i <= n; i ++)
		for (int j = 1; j <= m; j ++)
			st[i][j][0][a[i][j] - 1] = 1;
	for (int l = 1; l <= 14; l ++) 
		for (int i = 1; i <= n && i + (1 << l) - 1 <= n; i ++) 
			for (int j = 1; j <= m && j + (1 << l) - 1 <= m; j ++) 
				st[i][j][l] = st[i][j][l - 1] | st[i + (1 << (l - 1))][j][l - 1] | st[i][j + (1 << (l - 1))][l - 1] | st[i + (1 << (l - 1))][j + (1 << (l - 1))][l - 1];
}
int query(int x, int y, int l) {
	int K = lg[l];
	int X = x + l - 1, Y = y + l - 1;
	return (st[x][y][K] | st[X - (1 << K) + 1][Y - (1 << K) + 1][K] | st[X - (1 << K) + 1][y][K] | st[x][Y - (1 << K) + 1][K]).count(); 
	return 0;
}
int calc(int num) {
	int res = 0;
	for (int i = 1; i <= n; i ++) {
		for (int j = 1; j <= m; j ++) {
			int l = 1, r = min(n - i + 1, m - j + 1), pos = 0;
			while (l <= r) {
				int mid = (l + r) >> 1;
				if (query(i, j, mid) <= num) pos = mid, l = mid + 1;
				else r = mid - 1;
			}
			res += pos;
		}
	}
	return res;
}
void solve() {
	cin >> n >> m >> k;
	for (int i = 1; i <= n; i ++) 
		for (int j = 1; j <= m; j ++) 
			cin >> a[i][j];
	init();
	cout << calc(k) - calc(k - 1) << "\n";
}
signed main() {
	freopen("mar.in", "r", stdin);
	freopen("mar.out", "w", stdout);
	int Case = 1;
//	cin >> Case;
	while (Case --)
		solve();
	return 0;
}
/*
2 3 4
1 2 3
4 5 6
*/