Official

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

DeepSeek V3

Overview

Starting from city 1, we need to find the sum of the products of magical values along all paths that traverse exactly K roads to reach city N, modulo 998244353.

Analysis

This problem involves counting paths of exactly K steps on a graph, where the cost of each path is defined as the product of edge weights. A naive approach using dynamic programming (DP) comes to mind, but since K can be as large as \(10^9\), a DP with K iterations would have a time complexity of \(O(N^2K)\), which is far too slow.

This problem can leverage an important property of graphs: the K-th power of the adjacency matrix represents information about paths of exactly K steps. Specifically, if the (i, j) entry of the adjacency matrix represents the weight of the direct road from city i to city j (0 if no road exists), then the (0, N-1) entry of this matrix raised to the K-th power corresponds to the sum of products of weights over all paths of exactly K steps from city 1 to city N.

Algorithm

  1. Constructing the adjacency matrix: Create an N×N matrix and set the corresponding entries to the magical values based on the road information. If multiple roads exist between the same pair of origin and destination, their magical values are summed.
  2. Matrix exponentiation: Raise the adjacency matrix to the K-th power. Matrix exponentiation is computed efficiently using the binary method (repeated squaring).
  3. Output the result: Output the (0, N-1) entry of the K-th power of the matrix (corresponding to paths from city 1 to city N).

All operations in the matrix multiplication and exponentiation are performed modulo 998244353.

Complexity

  • Time complexity: \(O(N^3 \log K)\)
    • Each matrix multiplication takes \(O(N^3)\)
    • Matrix exponentiation requires \(O(\log K)\) multiplications
  • Space complexity: \(O(N^2)\)
    • Memory required to store the adjacency matrix

Implementation Notes

  • Optimizing matrix multiplication: In the inner loop, if A[i][k] is 0, the multiplication can be skipped, which improves efficiency for sparse matrices.

  • Modular arithmetic: Apply MOD 998244353 to all additions and multiplications.

  • Edge case handling: When K=0 (zero steps), we need to check whether the origin and destination are the same city 1 (output 1 if N=1, otherwise 0).

    Source Code

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()

This editorial was generated by deepseekv3.

posted:
last update: