G - Random Walk to Millionaire Editorial by en_translator
Consider a walk \(\pmb{v}\coloneqq (v_0, v_1, v_2, \ldots, v_K)\) resulting from \(K\) actions, where \(v_i\) is the vertex you move to after the \(i\)-th action.
When he passes \(\pmb{v}\), the total amount \(f(\pmb{v})\) of money that he obtains is expressed as \(f(\pmb{v}) = \displaystyle\sum_{\substack{i = 1, 2, \ldots, K \\ C_i = 1}} X_i^2\), where \(X_i\) is Takahashi’s Level after \(i\)-th action.
\(X_i\) equals the number of \(j\)’s such that \(j = 1, 2, \ldots, i\) and \(C_{v_j} = 0\), so \(X_i^2\) equals “the number of integer pairs \((x, y)\) such that \(1 \leq x, y \leq i\) and \(C_{v_x} = C_{v_y} = 0\).”
Therefore, let \(T\) be a set of all triples \((\pmb{w}, x, y)\) consisting of
- a walk \(\pmb{w} = (w_0, w_1, \ldots, w_l)\) and
- integers \(x\) and \(y\) such that \(1 \leq x, y \leq l\) and \(C_{w_x} = C_{w_y} = 0\);
then \(f(\pmb{v})\) equals:
the number of triples \((\pmb{w}, x, y) \in \mathcal{T}\) such that
- \(\pmb{w}\) is a prefix of \(\pmb {v}\) and
- the last vertex \(w_{\mathrm{last}}\) of \(\pmb{w}\) satisfies \(C_{w_{\mathrm{last}}} =1\).
Hence, the answer to this problem, i.e. the expected value \(E_{\pmb{v}}[f(\pmb{v})]\) of \(f(\pmb{v})\) when the walk \(\pmb{v}\) is chosen at random is:
the sum of “the probability that \((\pmb{w}, x, y)\) contributes to \(f(\pmb{v})\)” over all triples \((\pmb{w}, x, y) \in \mathcal{T}\).
When \(p(\pmb{w})\) denotes the probability that \((\pmb{w}, x, y)\) contributes to \(f(\pmb{v})\), it equals
\(E_{\pmb{v}}[f(\pmb{v})] = \displaystyle\sum_{\substack{(\pmb{w}, x, y) \in \mathcal{T} \\ C_{w_{\mathrm{last}}} = 1}} p(\pmb{w}) = \displaystyle\sum_{\substack{(\pmb{w}, x, y) \in \mathcal{T} \\ C_{w_{\mathrm{last}}} = 1}}\prod_{i = 0}^{l-1} \frac{1}{d(w_i)}\),
where \(d(w_i)\) denotes the degree of vertex \(w_i\).
We can evaluate it with a DP (Dynamic Programming). Specifically, we define the DP table \(\mathrm{dp}[i][j][a][b]\) by
the sum of \(p(\pmb{w})\) over all triples \((\pmb{w}, x, y)\) consisting of
- a walk \(\pmb{w} = (w_0, w_1, \ldots, w_i)\) of length \(i\) such that \(w_i = j\);
- an integer \(x\) such that
- \(1 \leq x \leq i\) and \(C_{w_x} = 0\) if \(a = 1\),
- \(x = -1\) if \(a = 0\); and
- an integer \(y\) such that
- \(1 \leq y \leq i\) and \(C_{w_y} = 0\) if \(b = 1\),
- \(y = -1\) if \(b = 0\).
By filling the table, we can find the answer as
\(E_{\pmb{v}}[f(\pmb{v})] = \displaystyle\sum_{\substack{(\pmb{w}, x, y) \in \mathcal{T} \\ C_{w_{\mathrm{last}}} = 1}} p(\pmb{w}) = \sum_{i = 1}^K \sum_{\substack{j = 1, 2, \ldots, N \\ C_j = 1}} \mathrm{dp}[i][j][1][1]\)
in a total of \(O(MK)\) time.
The following is a sample in C++ language.
#include <iostream>
#include <vector>
#include <atcoder/modint>
using namespace std;
using namespace atcoder;
typedef modint998244353 mint;
int n, m, k;
vector<int> G[3005];
int c[3005];
mint dp[3005][3005][2][2];
mint dinv[3005];
int main(void)
{
cin >> n >> m >> k;
int u, v;
for(int i = 1; i <= m; i++){
cin >> u >> v;
G[u].push_back(v), G[v].push_back(u);
}
for(int i = 1; i <= n; i++) cin >> c[i];
for(int i = 1; i <= n; i++) dinv[i] = mint(G[i].size()).inv();
dp[0][1][0][0] = 1;
for(int i = 0; i < k; i++){
for(int j = 1; j <= n; j++){
for(int a = 0; a < 2; a++){
for(int b = 0; b < 2; b++){
for(auto u : G[j]){
for(int na = a; na < 2; na++){
for(int nb = b; nb < 2; nb++){
if(c[u] == 1 && na+nb != a+b) continue;
dp[i+1][u][na][nb] += dp[i][j][a][b] * dinv[j];
}
}
}
}
}
}
}
mint ans = 0;
for(int i = 1; i <= k; i++) for(int j = 1; j <= n; j++) if(c[j]) ans += dp[i][j][1][1];
cout << ans.val() << endl;
return 0;
}
posted:
last update: