R - Walk atcoder educational dp


here is code…
mod = 10**9 + 7

N, K = map(int,input().split())
A = [list(map(int,input().split())) for _ in range(N)]
 
def dot(X, Y):
    x = len(X)
    y = len(Y[0])
    z = len(Y)
    C = [[0] * y for _ in range(x)]
    for i in range(x):
        for j in range(y):
            for k in range(z):
                C[i][j] += X[i][k] * Y[k][j]
                C[i][j] %= mod
    return C
 
B = [[0] * N for _ in range(N)]
for k in range(N):
    B[k][k] = 1
 
while K > 0:
    if K & 1:
        B = dot(B, A)
    A = dot(A, A)
    K >>= 1
 
ans = 0
for i in range(N):
    for j in range(N):
        ans += B[i][j]
        ans %= mod
 
print(ans)

The above code is working.How matrix exponenetian could give the ans.Please help me how to prove the logic is correct.

Matrix exponentiation is just a side effect.
Say, I know the number of ways to get from i to j for all i and j in length a and b. How will we find it for a+b?
let dp1[i][j] denote the number of ways to get from i to j in a steps and dp2[i][j] denote the number of ways to get from i to j in b steps. I must be somewhere after a steps. Let’s call that k. Now the number of ways to get from i to k in a steps is dp1[i][k] and the number of ways to get from k to j is dp2[k][j]. So the number of ways to go from i to j in a+b steps when I go to k after a steps is
dp1[i][k] \times dp2[k][j].
Therefore the number of ways to go to i from j after a+b steps is
\sum dp1[i][k] \times dp2[k][j] for all k.
Using this relation we can find the number of ways to go from i to j in x steps in O(log \ x) time.

11 Likes

thanks again.

Great explanation