Codeforces 253D Table with Letters - 2

問題

nxmのアルファベット小文字の書かれたグリッドがある。
このグリッドの部分長方形うち、四隅のアルファベットが同じで、含まれている'a'の個数がk個であるようなものの個数を求めよ。

制約条件

n, m≦400
k≦nm

方針

尺取法。

アルファベットを固定する。

行を上から見ていって、同じ行のj番目, k番目の文字が同じだったときに、v[j][k]にiを入れていくような処理の仕方を考える。

すると、j, k番目に同じ文字を見つけたとき、
今までの行のうち、どれに、j, k番目に同じ文字があるかは、
v[j][k]を舐めればわかる。

ここで、v[j][k]ごとに添え字idx[j][k]を記憶しておけば、
尺取法をすることができる。
すなわち、今の行でギリギリ長方形が作れるidx≦i+1行以降でギリギリ長方形が作れるidx
がなりたつので、前の行のときに保存したidxから、比較をはじめればよくって、
全体として(j, k)の組一つにつきO(n)回の比較で、全部の長方形の数を数えられる。

比較には、二次元累積和をあらかじめ計算しておけば一回あたりO(1)の時間しかかからない。

が、これだと最悪ケースでベクターに要素がO(nm^2)個入ってしまってMLE.
なので、ループの順番を少し変えて、
右側の文字を固定してやるようにする。
するとオーダーが変わらずに、メモリ消費量だけ減って通る。

ソースコード

int n, m, K;
char in[401][401];

vi v[400][26];
int idx[400][26], a[401][401];

int main(){
	
	freopen("input.txt", "r", stdin);
	freopen("output.txt", "w", stdout);
	scanf("%d%d%d", &n, &m, &K);
	rep(i, n){
		scanf("%s", in[i]);
		rep(j, m) a[i+1][j+1] = a[i+1][j] + a[i][j+1] - a[i][j] + (in[i][j] == 'a');
	}
	
	ll ans = 0;
	
	rep(r, m){
		rep(i, m) rep(j, 26){
			idx[i][j] = 0;
			v[i][j].clear();
		}
		rep(i, n){
			rep(j, r) if(in[i][j] == in[i][r]){
				
				vi &vv = v[j][in[i][j] - 'a'];
				int &ii = idx[j][in[i][j] - 'a'];
				int x = a[i+1][r+1] - a[i+1][j];
				
				while(ii < vv.size() && x - vv[ii] > K) ii++;
				
				ans += vv.size() - ii;
				vv.pb(a[i][r+1] - a[i][j]);
			}
		}
	}
	cout << ans << endl;
	
	return 0;
}