kmjp's blog

競技プログラミング参加記です

Codeforces #510 Div2 E. Vasya and Magic Matrix

なんか最近Div2ばっかりですね。
http://codeforces.com/contest/1042/problem/E

問題

整数が書かれた二次元グリッドと初期位置が与えられる。
現在位置から、自身の位置のマスより小さな値が書かれたマスに当確率で遷移する、という作業を最小値に到達するまで繰り返す。
移動距離の二乗和の期待値を求めよ。

解法

まずマス目を書かれた整数別に分類する。
距離の二乗和を計算したいので、X座標とY座標の総和および二乗和を移動元の到達確率の重みづけでとった平均値を持って、適宜移動先の到達確率と距離の二乗和の期待値を更新していけばよい。

int H,W;
int A[1010][1010];
int R,C;
map<int,vector<pair<int,int>>> M;
ll mo=998244353;

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>H>>W;
	FOR(y,H) FOR(x,W) cin>>A[y][x];
	
	cin>>R>>C;
	R--,C--;
	M[A[R][C]].push_back({C,R});
	FOR(y,H) FOR(x,W) if(A[y][x]<A[R][C]) M[A[y][x]].push_back({x,y});
	
	
	ll ret=0;
	int sum=0;
	ll P=1;
	int num=0;
	ll X2=C*C,Y2=R*R,X=C,Y=R;
	vector<vector<pair<int,int>>> step;
	FORR(m,M) if(m.first<A[R][C]) num+=m.second.size(), step.push_back(m.second);
	reverse(ALL(step));
	
	FORR(s,step) {
		FORR(v,s) {
			(ret+=((v.first*v.first+v.second*v.second+X2+Y2-2*v.first*X-2*v.second*Y)%mo+mo)%mo*modpow(num))%=mo;
		}
		X2=X2*(num-s.size())%mo;
		Y2=Y2*(num-s.size())%mo;
		X=X*(num-s.size())%mo;
		Y=Y*(num-s.size())%mo;
		FORR(v,s) {
			X2+=v.first*v.first;
			Y2+=v.second*v.second;
			X+=v.first;
			Y+=v.second;
		}
		X2=X2*modpow(num)%mo;
		Y2=Y2*modpow(num)%mo;
		X=X*modpow(num)%mo;
		Y=Y*modpow(num)%mo;
		num-=s.size();
	}
	
	cout<<ret%mo<<endl;
}

まとめ

地味に除算が多くて最初TLEした。