Official

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: