Problem p03177 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
#define ll long long int
#define M 1000000007

ll walk(vector<vector<ll>> &adj, ll n, ll K) {

  ll dp[K + 1][n][n];
  memset(dp, 0, sizeof(dp));
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++)
      dp[1][i][j] = adj[i][j];

  for (int k = 2; k <= K; k++) {
    for (int z = 0; z < n; z++)
      for (int i = 0; i < n; i++)
        for (int j = 0; j < n; j++) {
          // No. of paths from i to j in 'k' steps +=
          // No. of paths from i to z in (k-1) steps *
          // No. of paths from z to j in (k-1) steps
          dp[k][i][j] =
              (dp[k][i][j] + (dp[k - 1][i][z] * dp[k - 1][z][j]) % M) % M;
        }
  }

  ll ans = 0;
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++)
      ans = (ans + dp[K][i][j]) % M;
  return ans;
}

void multiply(vector<vector<ll>> &A, vector<vector<ll>> &B, ll n) {

  // Multiply the matrices 'A' and 'B' and store the result in 'A' itself
  vector<vector<ll>> C(n, vector<ll>(n));
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++) {
      C[i][j] = 0;
      for (int k = 0; k < n; k++)
        C[i][j] = (C[i][j] + (A[i][k] * B[k][j]) % M) % M;
    }

  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++)
      A[i][j] = C[i][j];
}

void matrixExponentiation(vector<vector<ll>> &A, ll n, ll k) {

  if (k <= 1)
    return;
  if (k % 2 == 0) {
    matrixExponentiation(A, n, k / 2);
    multiply(A, A, n);
  } else {
    vector<vector<ll>> tempMatrix(A);
    matrixExponentiation(A, n, k - 1);
    multiply(A, tempMatrix, n);
  }
}

ll walkOptimized(vector<vector<ll>> &adj, ll n, ll K) {
  // 'adj' contains for each cell (i, j) the no. of ways to
  // reach from 'i' to 'j' in  1 step

  // Suppose If we have a matrix 'A', in which each cell
  // (i, j) denotes the number of paths from i to j in exactly 4 steps
  // Then if we multiply it with 'adj', then after multiplying we will get
  // a matrix in which for each cell (i, j) it stores the no. of ways
  // to reach from 'i' to 'j' in exactly 5 steps

  // Hence we just need to multiply 'adj' K times, i.e.
  // we need to find adj^K, then in the end take the sum of all values
  // in that matrix because the paths can start or end at any vertex

  // To quickly find this exponentiation, we can use matrix exponentiation

  matrixExponentiation(adj, n, K);

  ll ans = 0;
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++)
      ans = (ans + adj[i][j]) % M;
  return ans;
}

int main() {
  ll n, k;
  cin >> n >> k;
  vector<vector<ll>> adj(n, vector<ll>(n));
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++)
      cin >> adj[i][j];
  ll ans = walk(adj, n, k);
  // ll ans = walkOptimized(adj, n, k);
  cout << ans;
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define ll long long int
#define M 1000000007

ll walk(vector<vector<ll>> &adj, ll n, ll K) {

  ll dp[K + 1][n][n];
  memset(dp, 0, sizeof(dp));
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++)
      dp[1][i][j] = adj[i][j];

  for (int k = 2; k <= K; k++) {
    for (int z = 0; z < n; z++)
      for (int i = 0; i < n; i++)
        for (int j = 0; j < n; j++) {
          // No. of paths from i to j in 'k' steps +=
          // No. of paths from i to z in (k-1) steps *
          // No. of paths from z to j in (k-1) steps
          dp[k][i][j] =
              (dp[k][i][j] + (dp[k - 1][i][z] * dp[k - 1][z][j]) % M) % M;
        }
  }

  ll ans = 0;
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++)
      ans = (ans + dp[K][i][j]) % M;
  return ans;
}

void multiply(vector<vector<ll>> &A, vector<vector<ll>> &B, ll n) {

  // Multiply the matrices 'A' and 'B' and store the result in 'A' itself
  vector<vector<ll>> C(n, vector<ll>(n));
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++) {
      C[i][j] = 0;
      for (int k = 0; k < n; k++)
        C[i][j] = (C[i][j] + (A[i][k] * B[k][j]) % M) % M;
    }

  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++)
      A[i][j] = C[i][j];
}

void matrixExponentiation(vector<vector<ll>> &A, ll n, ll k) {

  if (k <= 1)
    return;
  if (k % 2 == 0) {
    matrixExponentiation(A, n, k / 2);
    multiply(A, A, n);
  } else {
    vector<vector<ll>> tempMatrix(A);
    matrixExponentiation(A, n, k - 1);
    multiply(A, tempMatrix, n);
  }
}

ll walkOptimized(vector<vector<ll>> &adj, ll n, ll K) {
  // 'adj' contains for each cell (i, j) the no. of ways to
  // reach from 'i' to 'j' in  1 step

  // Suppose If we have a matrix 'A', in which each cell
  // (i, j) denotes the number of paths from i to j in exactly 4 steps
  // Then if we multiply it with 'adj', then after multiplying we will get
  // a matrix in which for each cell (i, j) it stores the no. of ways
  // to reach from 'i' to 'j' in exactly 5 steps

  // Hence we just need to multiply 'adj' K times, i.e.
  // we need to find adj^K, then in the end take the sum of all values
  // in that matrix because the paths can start or end at any vertex

  // To quickly find this exponentiation, we can use matrix exponentiation

  matrixExponentiation(adj, n, K);

  ll ans = 0;
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++)
      ans = (ans + adj[i][j]) % M;
  return ans;
}

int main() {
  ll n, k;
  cin >> n >> k;
  vector<vector<ll>> adj(n, vector<ll>(n));
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++)
      cin >> adj[i][j];
  // ll ans = walk(adj, n, k);
  ll ans = walkOptimized(adj, n, k);
  cout << ans;
  return 0;
}