Official

E - 魔法の道 / Magic Road Editorial by admin

DeepSeek V3

概要

街1から出発してちょうどK本の道を通って街Nに至るすべての経路について、通った道の魔力値の積の総和を998244353で割った余りを求める問題です。

考察

この問題は、グラフ上でちょうどKステップの経路を数え上げる問題ですが、各経路のコストが道の重みの積で定義されるという特徴があります。素朴なアプローチとして動的計画法(DP)が考えられますが、Kが最大10^9まで大きくなるため、K回の反復を行うDPでは計算量がO(N²K)となり、到底間に合いません。

この問題は、隣接行列のK乗がちょうどKステップの経路の情報を表現するというグラフの重要な性質を利用できます。具体的には、隣接行列の(i,j)成分が街iから街jへの直接の道の重み(存在しない場合は0)であるとき、この行列をK乗した行列の(0,N-1)成分が、街1から街NへのちょうどKステップの経路における重みの積の総和に対応します。

アルゴリズム

  1. 隣接行列の作成: N×Nの行列を作成し、各道の情報に基づいて対応する要素に魔力値を設定します。同じ出発地・到着地の組に対して複数の道が存在する場合は、それらの魔力値を加算します。
  2. 行列の累乗: 隣接行列をK乗します。行列の累乗はバイナリ法(繰り返し二乗法)を用いて効率的に計算します。
  3. 結果の出力: 行列のK乗結果の(0,N-1)成分(街1から街Nへの経路に対応)を出力します。

行列の乗算および累乗の計算において、すべての演算はMOD 998244353で行います。

計算量

  • 時間計算量: \(O(N^3 \log K)\)
    • 行列乗算1回あたり\(O(N^3)\)
    • 行列の累乗に\(O(\log K)\)回の乗算が必要
  • 空間計算量: \(O(N^2)\)
    • 隣接行列の保持に必要なメモリ

実装のポイント

  • 行列乗算の最適化: 内側のループで、A[i][k]が0の場合は乗算をスキップするようにすることで、疎な行列の場合に計算を効率化できます。

  • モジュラ演算: すべての加算・乗算でMOD 998244353を適用します。

  • エッジケースの処理: K=0の場合(ステップ数0)は、出発地と目的地が同じ街1かどうかで場合分けが必要です(N=1なら1、それ以外なら0)。

    ソースコード

MOD = 998244353

def mat_mult(A, B):
    n = len(A)
    m = len(B)
    p = len(B[0])
    C = [[0] * p for _ in range(n)]
    for i in range(n):
        for k in range(m):
            if A[i][k]:
                for j in range(p):
                    C[i][j] = (C[i][j] + A[i][k] * B[k][j]) % MOD
    return C

def mat_pow(matrix, power):
    n = len(matrix)
    result = [[1 if i == j else 0 for j in range(n)] for i in range(n)]
    base = matrix
    while power:
        if power & 1:
            result = mat_mult(result, base)
        base = mat_mult(base, base)
        power //= 2
    return result

def main():
    import sys
    data = sys.stdin.read().split()
    if not data:
        print(0)
        return
        
    it = iter(data)
    N = int(next(it)); M = int(next(it)); K = int(next(it))
    
    adj = [[0] * N for _ in range(N)]
    for i in range(M):
        u = int(next(it)) - 1
        v = int(next(it)) - 1
        w = int(next(it))
        adj[u][v] = (adj[u][v] + w) % MOD
        
    if K == 0:
        if N == 1:
            print(1)
        else:
            print(0)
        return
        
    T = mat_pow(adj, K)
    print(T[0][N-1] % MOD)

if __name__ == '__main__':
    main()

この解説は deepseekv3 によって生成されました。

posted:
last update: