CF510 E - Vasya and Magic Matrix

問題

h*wの行列aがあり、最初は(r,c)にチップが置いてある
Vasyaはランダムな「今チップが置いてあるマスに書かれた数より小さい数が書いてあるマス」にチップを動かすという操作を、チップを動かせなくなるまで続ける
チップを動かす度に、動かす前のマスと動かした後のマス間のユークリッド距離の2乗をスコアに加える
最終的なスコアの期待値を求めよ

1 <= h,w <= 1000
1 <= a[i][j] <= 10^9

解法

マス(R,C)から始まるときの期待値E(R,C)は、(R,C)に書かれた数未満の数が書かれたマスの集合を(r1,c1), (r2,c2), ..., (rk,ck)とすると、
Σ(E(Ri,Ci) + (R-ri)^2 + (C-ci)^2) / k
= (ΣE(Ri,Ci) + Σri^2 + Σci^2 - 2RΣri - 2CΣci + kR^2 + kC^2) / k
となる
これはマスに書かれた数ごとにマスの座標をmapとかでグループ分けし、数が小さい方から順に一個ずつ見ていき、「期待値・x座標・x座標2乗・y座標・y座標2乗」それぞれの和をとっていけば高速に求められる

また期待値は分数で表す必要があり、そこは二分累乗とフェルマーの小定理 PQ^(-1) ≡ PQ^(m-2) (mod m) を用いればいい

トータルでO(HWlog(HW)log(maxa))くらい？

#define int long long
const int mod = 998244353

int h, w, a[1010][1010], r, c, x, xx, y, yy, e, t;
map<int,vector<pair<int,int> > > m;

signed main() {
	cin >> h >> w;
	rep(i,h) rep(j,w) {
		cin >> a[i][j];
		m[a[i][j]].push_back({i,j});
	}
	cin >> r >> c;
	r--, c--;
	each(k,m) {
		if (k.first==a[r][c]) break;
		int xs = 0, xxs = 0, ys = 0, yys = 0, es = 0, ts = 0;
		each(i,k.second) {
			int xi = i.first, yi = i.second;
			xs += xi, xxs += xi*xi, ys += yi, yys += yi*yi, ts++;
			if (t!=0) {
				int p = e + xi*xi*t + yi*yi*t + xx + yy - 2*xi*x - 2*yi*y;
				es += p %mod * mop(t,mod-2) %mod;
			}
		}
		x += xs, xx += xxs, y += ys, yy += yys, e += es, t += ts;
	}
	int p = e + r*r*t + c*c*t + xx + yy - 2*r*x - 2*c*y;
	cout << p %mod * mop(t,mod-2) %mod << endl;
}

この小さい方から和を取っていくって発想に至らなかったのは反省だなぁ