Official

E - 粒子の拡散シミュレーション / Particle Diffusion Simulation Editorial by admin

gpt-5.5-high

概要

粒子の「場所」と「電荷」を状態として持つ拡散を、\(T\) 回後まで高速に計算する問題です。
\(T\) が最大 \(10^{18}\) と非常に大きいので、行列累乗を使って解きます。

考察

各容器について、正の粒子量 \(x_i^+\) と負の粒子量 \(x_i^-\) をそのまま管理すると、状態数は \(2N\) 個になります。

しかし、ここで次の 2 つの量を考えると問題が簡単になります。

  • 合計量:\(a_i = x_i^+ + x_i^-\)
  • 差分量:\(b_i = x_i^+ - x_i^-\)

このとき、正の粒子量は後で

\[ x_i^+ = \frac{a_i + b_i}{2} \]

として復元できます。

合計量について

電荷が反転しても、粒子の総量は変わりません。

したがって、容器 \(i\) から容器 \(j\) へ割合 \(P_{i,j}\) で移動するなら、合計量は単純に

\[ a'_j = \sum_i P_{i,j} a_i \]

となります。

差分量について

差分量 \(b_i = x_i^+ - x_i^-\) は、電荷が反転すると符号が反転します。

例えば、容器 \(i\) から \(j\) への遷移で電荷が反転する場合、

  • 正の粒子は負になる
  • 負の粒子は正になる

ので、到着後の差分は

\[ x_i^- P_{i,j} - x_i^+ P_{i,j} = - (x_i^+ - x_i^-) P_{i,j} = -b_i P_{i,j} \]

となります。

反転しない場合はそのまま

\[ b_i P_{i,j} \]

です。

よって、差分量は

\[ b'_j = \sum_i c_{i,j} P_{i,j} b_i \]

と書けます。ただし、

\[ c_{i,j} = \begin{cases} 1 & \text{反転しない場合} \\ -1 & \text{反転する場合} \end{cases} \]

です。

素朴な方法では間に合わない理由

1 ステップずつシミュレーションすると、各ステップで全ての遷移を計算する必要があり、計算量は \(O(TN^2)\) になります。

\(T\) は最大 \(10^{18}\) なので、これは到底間に合いません。

そこで、1 ステップの更新を行列で表し、行列累乗によって \(T\) ステップ分を一気に計算します。

アルゴリズム

列ベクトルとして状態を持つことにします。

合計量用の行列 \(A\) を次のように定義します。

\[ A_{j,i} = P_{i,j} \]

これは「容器 \(i\) から容器 \(j\) へ移動する割合」を表します。

すると、合計量ベクトル \(a\) は

\[ a' = A a \]

で更新できます。

次に、差分量用の行列 \(B\) を定義します。

\[ B_{j,i} = \begin{cases} P_{i,j} & \text{反転しない場合} \\ -P_{i,j} & \text{反転する場合} \end{cases} \]

すると、差分量ベクトル \(b\) は

\[ b' = B b \]

で更新できます。

初期状態では、容器 \(S\) に正の粒子が \(1\) だけあるので、

\[ a_S = 1,\quad b_S = 1 \]

であり、それ以外は \(0\) です。

\(T\) ステップ後は、

\[ a^{(T)} = A^T a^{(0)} \]

\[ b^{(T)} = B^T b^{(0)} \]

です。

最後に、容器 \(G\) の正の粒子量は

\[ x_G^+ = \frac{a_G^{(T)} + b_G^{(T)}}{2} \]

で求められます。

分数の扱い

遷移確率は

\[ P_{i,j} = \frac{R_{i,j}}{\sum_k R_{i,k}} \]

です。

答えは \(10^9+7\) での値を出力するので、割り算は逆元で処理します。

つまり、

\[ P_{i,j} \]

はプログラム上では

\[ R_{i,j} \times \left(\sum_k R_{i,k}\right)^{-1} \pmod{10^9+7} \]

として計算します。

\(10^9+7\) は素数なので、逆元はフェルマーの小定理により

\[ x^{-1} \equiv x^{MOD-2} \pmod{MOD} \]

で求められます。

また、最後の \(\frac{1}{2}\) も逆元を使って掛け算します。

\[ \frac{1}{2} \equiv 500000004 \pmod{10^9+7} \]

です。

計算量

  • 時間計算量: \(O(N^3 \log T)\)
  • 空間計算量: \(O(N^2)\)

行列は合計量用と差分量用の 2 種類を持ちますが、定数倍なので計算量は \(O(N^3 \log T)\) です。

実装のポイント

行列の添字の向きに注意が必要です。

容器 \(i\) から容器 \(j\) への遷移確率を、行列では

\[ A_{j,i} \]

に入れています。

これは列ベクトル \(v\) に対して

\[ v' = A v \]

と更新するためです。

コード中では次のようにしています。

total_mat[j][i] = val

差分量用の行列では、反転ルールがある場合だけ符号を反転させます。

if flip[i][j]:
    diff_mat[j][i] = MOD - val
else:
    diff_mat[j][i] = val

負の値 \(-val\) は、mod 上では MOD - val として表します。

最後に、

ans = (total_vec[G] + diff_vec[G]) * INV2 % MOD

として、合計量と差分量から正の粒子量を復元します。

ソースコード

import sys

MOD = 1000000007
INV2 = (MOD + 1) // 2

def mat_mul(A, B, n):
    res = []
    mod = MOD
    for i in range(n):
        row_res = [0] * n
        Ai = A[i]
        for k in range(n):
            a = Ai[k]
            if a:
                Bk = B[k]
                for j in range(n):
                    row_res[j] += a * Bk[j]
        for j in range(n):
            row_res[j] %= mod
        res.append(row_res)
    return res

def mat_vec_mul(A, v, n):
    res = [0] * n
    mod = MOD
    for i in range(n):
        s = 0
        Ai = A[i]
        for j in range(n):
            s += Ai[j] * v[j]
        res[i] = s % mod
    return res

def main():
    data = list(map(int, sys.stdin.buffer.read().split()))
    idx = 0

    N = data[idx]
    M = data[idx + 1]
    T = data[idx + 2]
    idx += 3

    S = data[idx] - 1
    G = data[idx + 1] - 1
    idx += 2

    R = []
    row_sums = []
    for _ in range(N):
        row = data[idx:idx + N]
        idx += N
        R.append(row)
        row_sums.append(sum(row))

    flip = [[False] * N for _ in range(N)]
    for _ in range(M):
        u = data[idx] - 1
        v = data[idx + 1] - 1
        idx += 2
        flip[u][v] = True

    total_mat = [[0] * N for _ in range(N)]
    diff_mat = [[0] * N for _ in range(N)]

    for i in range(N):
        inv_sum = pow(row_sums[i], MOD - 2, MOD)
        Ri = R[i]
        for j in range(N):
            val = Ri[j] * inv_sum % MOD
            total_mat[j][i] = val
            if flip[i][j]:
                diff_mat[j][i] = MOD - val if val else 0
            else:
                diff_mat[j][i] = val

    total_vec = [0] * N
    diff_vec = [0] * N
    total_vec[S] = 1
    diff_vec[S] = 1

    t = T
    while t:
        if t & 1:
            total_vec = mat_vec_mul(total_mat, total_vec, N)
            diff_vec = mat_vec_mul(diff_mat, diff_vec, N)
        t >>= 1
        if t:
            total_mat = mat_mul(total_mat, total_mat, N)
            diff_mat = mat_mul(diff_mat, diff_mat, N)

    ans = (total_vec[G] + diff_vec[G]) * INV2 % MOD
    print(ans)

if __name__ == "__main__":
    main()

この解説は gpt-5.5-high によって生成されました。