E - 粒子の拡散シミュレーション / Particle Diffusion Simulation Editorial by admin
claude4.8opus-high概要
\(N\) 個の容器の間を粒子が遷移し、特定の経路では電荷が反転するというシミュレーションを \(T\) ステップ行い、最終的に容器 \(G\) に存在する正電荷粒子の量を求める問題です。各容器が「正」「負」の \(2\) 状態を持つため状態数は \(2N\) 個であり、\(1\) ステップの更新が線形変換であることに着目して行列累乗で高速に計算します。
考察
1ステップの更新は線形変換
問題文の更新式をよく見ると、更新後の値 \(y_j^{+}, y_j^{-}\) はすべて更新前の値 \(x_i^{+}, x_i^{-}\) の一次結合(定数倍の和)になっています。
\[y_j^{+} = \sum_{i=1}^{N} \Bigl( (1 - F_{i,j}) x_i^{+} P_{i,j} + F_{i,j} x_i^{-} P_{i,j} \Bigr)\]
つまり、すべての容器の状態をまとめた長さ \(2N\) のベクトル
\[\mathbf{x} = (x_1^{+}, \dots, x_N^{+}, x_1^{-}, \dots, x_N^{-})\]
を考えると、\(1\) ステップの更新は、ある \(2N \times 2N\) 行列 \(A\) を用いて \(\mathbf{x}_{\text{new}} = A \mathbf{x}\) と表せます。
素朴なアプローチの問題点
\(T \leq 10^{18}\) なので、\(1\) ステップずつ愚直にシミュレーションすると到底間に合いません。しかし更新が線形変換であることから、\(T\) ステップ後の状態は
\[\mathbf{x}_T = A^T \mathbf{x}_0\]
で求まります。\(A^T\) は繰り返し二乗法(行列累乗)で \(O(\log T)\) 回の行列積で計算できます。
有理数の扱い
確率 \(P_{i,j} = R_{i,j} / (\sum_{j'} R_{i,j'})\) は有理数なので、\(10^9+7\) を法としたモジュラ逆元で割り算を表現します。行和は制約より \(1\) 以上なので逆元が存在し、問題の保証より最終的な答えの分母も \(10^9+7\) と互いに素です。
アルゴリズム
状態の番号付け
長さ \(2N\) のベクトルの添字を次のように割り当てます。
- インデックス \(i\)(\(0 \le i < N\)): 容器 \(i\) の正電荷の量
- インデックス \(N+i\): 容器 \(i\) の負電荷の量
行列 \(A\) の構築
各組 \((i, j)\) について \(P_{i,j} > 0\) のとき、遷移を行列に書き込みます。
- 反転なし(\(F_{i,j} = 0\))の場合: 正→正、負→負
\[A[\text{pos}_j][\text{pos}_i] \mathrel{+}= P_{i,j}, \quad A[\text{neg}_j][\text{neg}_i] \mathrel{+}= P_{i,j}\]
- 反転あり(\(F_{i,j} = 1\))の場合: 正→負、負→正
\[A[\text{neg}_j][\text{pos}_i] \mathrel{+}= P_{i,j}, \quad A[\text{pos}_j][\text{neg}_i] \mathrel{+}= P_{i,j}\]
同じ位置に複数の遷移が足し合わさることがあるため、代入ではなく加算で蓄積します。
初期ベクトルと答えの取り出し
初期状態は「容器 \(S\) の正電荷が \(1\)、他はすべて \(0\)」なので、\(\mathbf{x}_0\) は \(\text{pos}_S\) 成分のみ \(1\) の基底ベクトルです。
求める答えは \(\mathbf{x}_T = A^T \mathbf{x}_0\) の \(\text{pos}_G\) 成分です。これは初期ベクトルが基底ベクトルなので、
\[A^T \text{ の } (\text{pos}_G, \text{pos}_S) \text{ 成分}\]
そのものになります。よって行列 \(A^T\) さえ計算すれば、その \(1\) 成分を出力するだけで済みます。
計算量
行列のサイズは \(2N \times 2N\) で、\(N \le 30\) より最大 \(60 \times 60\) です。
- 行列積 \(1\) 回: \(O((2N)^3) = O(N^3)\)
- 繰り返し二乗法: \(O(\log T)\) 回の行列積
したがって
- 時間計算量: \(O(N^3 \log T)\)
- 空間計算量: \(O(N^2)\)
\(N = 30\)、\(T = 10^{18}\) でも \(60^3 \times 60 \approx 1.3 \times 10^7\) 程度の演算で、十分高速です。
実装のポイント
モジュラ逆元: 行和 \(s\) の逆元は
pow(s % mod, mod-2, mod)(フェルマーの小定理)で求めます。\(P_{i,j} = R_{i,j} \cdot s^{-1} \bmod p\) として確率を前計算しておくと簡潔です。加算での蓄積: 行列 \(A\) を作る際、\(P_{i,i} > 0\)(その場に留まる粒子)など複数の寄与が同じ成分に集まることがあるため、
A[...] = (A[...] + pij) % modのように足し込みます。答えの直接取得: 初期ベクトルが基底ベクトルであることを利用し、\(A^T\) の \((\text{pos}_G, \text{pos}_S)\) 成分を直接出力すれば、行列とベクトルの積を別途計算する必要がありません。
反転ルールの方向: 反転は有向なので、\(F[u][v]\) のみを \(1\) にし、\(F[v][u]\) は触れないように注意します。
ソースコード
import sys
def main():
input_data = sys.stdin.buffer.read().split()
idx = 0
N = int(input_data[idx]); idx+=1
M = int(input_data[idx]); idx+=1
T = int(input_data[idx]); idx+=1
S = int(input_data[idx]); idx+=1
G = int(input_data[idx]); idx+=1
mod = 10**9+7
R = [[0]*N for _ in range(N)]
for i in range(N):
for j in range(N):
R[i][j] = int(input_data[idx]); idx+=1
F = [[0]*N for _ in range(N)]
for _ in range(M):
u = int(input_data[idx])-1; idx+=1
v = int(input_data[idx])-1; idx+=1
F[u][v] = 1
# row sums and inverse
rowinv = [0]*N
for i in range(N):
s = sum(R[i])
rowinv[i] = pow(s % mod, mod-2, mod)
# P[i][j] = R[i][j]*rowinv[i] mod
P = [[ (R[i][j]%mod)*rowinv[i] % mod for j in range(N)] for i in range(N)]
size = 2*N
# build matrix A : new = A * old
# pos index = i, neg index = N+i
A = [[0]*size for _ in range(size)]
for i in range(N):
for j in range(N):
pij = P[i][j]
if pij == 0:
continue
fij = F[i][j]
pj = j
nj = N+j
pi = i
ni = N+i
if fij:
# pos_j gets from neg_i ; neg_j gets from pos_i
A[pj][ni] = (A[pj][ni] + pij) % mod
A[nj][pi] = (A[nj][pi] + pij) % mod
else:
A[pj][pi] = (A[pj][pi] + pij) % mod
A[nj][ni] = (A[nj][ni] + pij) % mod
def matmul(X, Y):
n = size
Z = [[0]*n for _ in range(n)]
for i in range(n):
Xi = X[i]
Zi = Z[i]
for k in range(n):
a = Xi[k]
if a:
Yk = Y[k]
for j in range(n):
Zi[j] = (Zi[j] + a*Yk[j]) % mod
return Z
# identity
res = [[1 if i==j else 0 for j in range(size)] for i in range(size)]
base = A
e = T
while e:
if e & 1:
res = matmul(res, base)
e >>= 1
if e:
base = matmul(base, base)
# answer = res[pos_G][pos_S]
posG = G-1
posS = S-1
print(res[posG][posS] % mod)
main()
この解説は claude4.8opus-high によって生成されました。