AtCoder ABC #277 (大和証券プログラミングコンテスト2022 Autumn) : G - Random Walk to Millionaire

これはよくあるタイプの問題かな…。
https://atcoder.jp/contests/abc277/tasks/abc277_g

問題

N頂点からなる連結無向グラフが与えられる。
また、各点は0/1いずれかの状態を持つ。

プレイヤーは、初期状態で0番の点におり、レベルは0とする。
以下K回以下の手順を繰りかえす。

• 現在いる頂点の隣接点のうち、いずれかを等確率で選択し、移動する。
  • 状態0の頂点に移動した場合、プレイヤーのレベルが1増える。
  • 状態1の頂点に移動した場合、プレイヤーのレベルの2乗のスコアを得る。

K回繰り返したときに得られる総スコアの期待値を求めよ。

解法

以下を状態として管理しよう。

• P(k,v) := k回移動して頂点vにいる確率
• L(k,v) := k回移動して頂点vにいるときのレベルの期待値*P(k,v)
• X(k,v) := k回移動して頂点vにいるときのレベルの二乗和の期待値*P(k,v)

k回目の移動では、移動元vから移動先wに対し、P(k-1,v)→P(k,w)、L(k-1,v)→L(k,w)、S(k-1,v)→S(k,w)の寄与を計算する。
その後、各頂点、移動先の状態に応じた処理を行う。

• 頂点vの状態が0の場合
  • レベルが1増えないといけないので、
    • L(k,v)にP(k,v)を加算する
    • S(k,v)に2*L(k,v)+P(k,v)を加算する
• 頂点vの状態が1の場合
  • 解にS(k,v)を計上する。

int N,M,K;
vector<int> E[3030];
int C[3030];
const ll mo=998244353;
ll FP[3030];
ll FL[3030];
ll FS[3030];
ll TP[3030];
ll TL[3030];
ll TS[3030];

const int NUM_=4001;
static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	inv[1]=fact[0]=factr[0]=1;
	for (int i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
	for (int i=1;i<=NUM_;++i) fact[i]=fact[i-1]*i%mo, factr[i]=factr[i-1]*inv[i]%mo;
	cin>>N>>M>>K;
	FOR(i,M) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	FOR(i,N) cin>>C[i];
	
	FP[0]=1;
	ll sum=0;
	while(K--) {
		ZERO(TP);
		ZERO(TL);
		ZERO(TS);
		
		FOR(i,N) {
			FORR(e,E[i]) {
				(TP[e]+=FP[i]*inv[E[i].size()])%=mo;
				(TL[e]+=FL[i]*inv[E[i].size()])%=mo;
				(TS[e]+=FS[i]*inv[E[i].size()])%=mo;
			}
		}
		
		FOR(i,N) {
			if(C[i]==0) {
				(TS[i]+=2*TL[i]+TP[i])%=mo;
				(TL[i]+=TP[i])%=mo;
			}
			else {
				(sum+=TS[i])%=mo;
			}
			FP[i]=TP[i];
			FL[i]=TL[i];
			FS[i]=TS[i];
		}
	}
	cout<<sum<<endl;
	
}

まとめ

あんまり迷わず解けた。