公式

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

GLM 5.2 (High)

概要

正負の電荷を持つ粒子が有向グラフ上を遷移するシミュレーションにおいて、\(T\) ステップ後の特定の容器にある正の電荷の粒子の量を求める問題です。\(T\) が非常に大きいため、行列累乗を用いて高速に計算します。

考察

粒子の遷移は、各ステップでの正の粒子・負の粒子の量を使って線形計算で表すことができます。素朴に \(T\) 回のシミュレーションを行うと \(O(NT)\) の計算量となり、\(T \leq 10^{18}\) であるためTLE(実行時間超過)になります。

そこで、遷移を行列で表すことを考えます。各容器 \(i\) について「正の粒子の量」と「負の粒子の量」の2つの状態を持たせるため、全体で \(2N\) 個の状態を管理します。 容器 \(i\) から容器 \(j\) への遷移を考えます。反転ルールがある場合(\(F_{i,j} = 1\))、正の粒子は負として、負の粒子は正として到着します。反転ルールがない場合(\(F_{i,j} = 0\))は、そのままの電荷で到着します。

これを \(2N \times 2N\) の遷移行列 \(A\) としてまとめることで、1ステップの更新を「状態ベクトルに行列 \(A\) を掛ける操作」として表現できます。\(T\) ステップ後の状態は行列 \(A\) を \(T\) 乗したものを掛けることで得られます。行列累乗を用いれば、\(O((2N)^3 \log T)\) で高速に計算できます。

アルゴリズム

  1. 状態の定義: 容器 \(i\)(\(0 \le i < N\))の正の粒子の量をインデックス \(2i\)、負の粒子の量をインデックス \(2i+1\) として扱います。
  2. 遷移行列の構築: 容器 \(i\) から容器 \(j\) への遷移割合 \(P_{i,j}\) と反転フラグ \(F_{i,j}\) を用いて、遷移行列 \(A\) の要素を次のように設定します。
    • \(A_{2j, 2i} = (1 - F_{i,j}) P_{i,j}\)
    • \(A_{2j, 2i+1} = F_{i,j} P_{i,j}\)
    • \(A_{2j+1, 2i} = F_{i,j} P_{i,j}\)
    • \(A_{2j+1, 2i+1} = (1 - F_{i,j}) P_{i,j}\) なお、\(P_{i,j}\) は \(R_{i,j}\) を行和で割ることで求まりますが、剰余算の世界では行和のモジュラ逆元を掛けることで計算します。
  3. 行列累乗: 初期状態は「容器 \(S\) の正の粒子が \(1\)」なので、状態ベクトルの \(2(S-1)\) 番目の要素が \(1\)、他は \(0\) です。よって、求める \(T\) ステップ後の容器 \(G\) の正の粒子の量は、行列 \(A^T\) の \((2(G-1), 2(S-1))\) 成分に等しくなります。
  4. 行列累乗の計算: 繰り返し二乗法を用いて \(A^T\) を計算し、該当成分を出力します。

計算量

  • 時間計算量: 行列のサイズを \(n = 2N\) とすると、行列積1回あたり \(O(n^3)\)、行列累乗で \(O(\log T)\) 回の行列積を行うため、\(O(N^3 \log T)\) となります。\(N \le 30\) なので十分高速です。
  • 空間計算量: 行列のサイズが \(2N \times 2N\) なので、\(O(N^2)\) となります。

実装のポイント

  • 割り算を余り付きで計算する際、フェルマーの小定理を用いて「行和の逆元」を掛けることで処理します。pow(s, MOD - 2, MOD) を利用します。

  • 行列積の内側のループでは、要素が \(0\) の場合の計算をスキップすることで定数倍を軽減しています。

  • 状態のインデックス(\(2i, 2i+1\) など)を混乱しないように注意して実装します。

    ソースコード

import sys

def solve():
    input = sys.stdin.readline
    N, M, T = map(int, input().split())
    S, G = map(int, input().split())
    R = [list(map(int, input().split())) for _ in range(N)]
    F = [[0] * N for _ in range(N)]
    for _ in range(M):
        u, v = map(int, input().split())
        F[u-1][v-1] = 1
    
    MOD = 10**9 + 7
    invS = [0] * N
    for i in range(N):
        s = sum(R[i]) % MOD
        invS[i] = pow(s, MOD - 2, MOD)
        
    n = 2 * N
    A = [[0] * n for _ in range(n)]
    for i in range(N):
        for j in range(N):
            p = R[i][j] * invS[i] % MOD
            f = F[i][j]
            # y_j^+ = (1-F) * x_i^+ * p + F * x_i^- * p
            A[2*j][2*i] = (1 - f) * p % MOD
            A[2*j][2*i+1] = f * p % MOD
            # y_j^- = F * x_i^+ * p + (1-F) * x_i^- * p
            A[2*j+1][2*i] = f * p % MOD
            A[2*j+1][2*i+1] = (1 - f) * p % MOD
            
    def mat_mul(A, B, n):
        C = [[0] * n for _ in range(n)]
        for i in range(n):
            Ai = A[i]
            Ci = C[i]
            for k in range(n):
                aik = Ai[k]
                if aik:
                    Bk = B[k]
                    for j in range(n):
                        Ci[j] = (Ci[j] + aik * Bk[j]) % MOD
        return C

    def mat_pow(A, p, n):
        res = [[0] * n for _ in range(n)]
        for i in range(n):
            res[i][i] = 1
        while p > 0:
            if p & 1:
                res = mat_mul(res, A, n)
            A = mat_mul(A, A, n)
            p >>= 1
        return res
    
    res = mat_pow(A, T, n)
    ans = res[2*(G-1)][2*(S-1)]
    print(ans)

if __name__ == '__main__':
    solve()

この解説は or-glm5.2-high によって生成されました。