G - Random Walk to Millionaire Editorial by kyopro_friends


この問題は、俗に積の和典型と呼ばれる手法で解くことができます。

気持ち

積の和典型とは、「『何かに関する積の和はいくらか?』という問題を『何かからそれぞれ1つずつ選ぶ方法は何通りか?』と読み替える」というものです。

期待値というのはほぼ和のようなものなので、今回の問題は「レベルの \(2\) 乗の和はいくらか?」というようなものです。よって、レベルの \(2\) 乗の部分をうまく何か数えるものに置き換えます。

解法

移動による \(C_v\) の列を \(X\) とします。すなわち、\(i\) 回目の移動後に頂点 \(V_i\) にいたとき、\(X_i=C_{V_i}\) と定めます。

レベルの2乗というのは、「\(X_i=0\) である \(i\) の個数の \(2\) 乗」であり、これは「(\(X_i=0\) である相異なる \(i\)\(2\) つ選ぶ方法) \(\times 2 +\) (\(X_i=0\) である \(i\)\(1\) つ選ぶ方法)」に等しいです。

従って、「列 \(X\)\((0,0,1)\) 及び \((0,1)\) を(連続とは限らない)部分列として期待値で何度含むか?」を求めることができれば元の問題を解くことができ、これはDPにより容易に求めることができます。計算量は \(O(K(N+M))\) です。

#include<bits/stdc++.h>
#include<atcoder/modint>
using namespace std;
using mint=atcoder::modint998244353;

int main(){
	int n,m,k;
	cin >> n >> m >> k;
	vector<vector<int>>G(n);
	for(int i=0;i<m;i++){
		int x,y;
		cin >> x >> y;
		x--,y--;
		G[x].push_back(y);
		G[y].push_back(x);
	}
	vector<int>c(n);
	for(int i=0;i<n;i++)cin >> c[i];
	
	vector<mint>coef(n);
	for(int v=0;v<n;v++)if(G[v].size())coef[v]=mint(1)/G[v].size();
	
	vector<vector<array<mint,3>>>dp(k+1,vector<array<mint,3>>(n));
	//dp[i][v][cnt]=i回目の移動後に頂点vにいて0をcnt個決めた
	dp[0][0][0]=1;
	//配るDP
	for(int i=0;i<k;i++)for(int v=0;v<n;v++)for(auto vv:G[v]){
		for(int cnt=0;cnt<3;cnt++){
			dp[i+1][vv][cnt]+=dp[i][v][cnt]*coef[v];
			if(cnt<2&&c[vv]==0)dp[i+1][vv][cnt+1]+=dp[i][v][cnt]*coef[v];
		}
	}
	
	mint ans=0;
	for(int i=1;i<=k;i++)for(int v=0;v<n;v++)if(c[v]==1){
		ans+=dp[i][v][1];
		ans+=dp[i][v][2]*2;
	}
	cout << ans.val() << endl;
}

おまけ

積の和典型は、\(0,1\) を値にとる確率変数を適切に定めることで、ad-hocな言い換えをせず数式をいじるだけで同じ結果を得ることができる場合がほとんどです。今回の例では、\(i\) 回目の移動後に \(C_v=0\) である頂点にいれば \(1\) 、そうでなければ \(0\) をとる確率変数 \(X_i\) を定めることで、求めるものが \(2\sum_{i<j<k}E[X_iX_j(1-X_k)]+\sum_{i<j}E[X_i(1-X_j)]\) となることがわかり、上で述べたものと全く同じDP解法を得ることができます。

posted:
last update: