AtCoder Educational DP Contest R問題解いてみた

今回の問題

問題文

$N$ 頂点の単純有向グラフ $G$ があります。 頂点には $1, 2, \ldots, N$ と番号が振られています。

各 $i, j$ ($1 \leq i, j \leq N$) について、頂点 $i$ から $j$ への有向辺の有無が整数 $a_{i,j}$ によって与えられます。

• $a_{i,j} = 1$ ならば頂点 $i$ から $j$ への有向辺が存在し、
• $a_{i,j} = 0$ ならば存在しません。

$G$ の長さ $K$ の有向パスは何通りでしょうか？
$10^9 + 7$ で割った余りを求めてください。
ただし、同じ辺を複数回通るような有向パスも数えるものとします。

制約

• 入力はすべて整数である。
• $1 \leq N \leq 50$
• $1 \leq K \leq 10^{18}$
• $a_{i,j}$ は $0$ または $1$ である。
• $a_{i,i} = 0$

回答

#include <bits/stdc++.h>
using namespace std;

const int MOD = 1e9 + 7;

// 行列を表す型 (N x N の2次元配列)
typedef vector<vector<long long>> Matrix;

// 行列の掛け算 C = A * B
Matrix multiply(const Matrix &A, const Matrix &B) {
    int n = A.size();
    Matrix C(n, vector<long long>(n, 0));
    
    for (int i = 0; i < n; i++) {
        for (int k = 0; k < n; k++) {
            // A[i][k] が 0 なら計算しても意味がないので枝刈り（少し高速化）
            if (A[i][k] == 0) continue;
            
            for (int j = 0; j < n; j++) {
                C[i][j] = (C[i][j] + A[i][k] * B[k][j]) % MOD;
            }
        }
    }
    return C;
}

// 行列の累乗 Res = A^k (繰り返し二乗法)
Matrix power(Matrix A, long long k) {
    int n = A.size();
    Matrix Res(n, vector<long long>(n, 0));
    
    // 単位行列で初期化 (Res[i][i] = 1, 他は0)
    for (int i = 0; i < n; i++) Res[i][i] = 1;
    
    while (k > 0) {
        // ビットが立っていれば掛ける
        if (k & 1) Res = multiply(Res, A);
        
        // A を二乗する
        A = multiply(A, A);
        k >>= 1;
    }
    return Res;
}

int main() {
    int N;
    long long K;
    cin >> N >> K;
    
    Matrix A(N, vector<long long>(N));
    for(int i = 0; i < N; i++) {
        for(int j = 0; j < N; j++) {
            cin >> A[i][j];
        }
    }

    // 行列累乗で A^K を求める
    Matrix AnsMat = power(A, K);

    // 全成分の和が答え
    long long ans = 0;
    for(int i = 0; i < N; i++) {
        for(int j = 0; j < N; j++) {
            ans = (ans + AnsMat[i][j]) % MOD;
        }
    }

    cout << ans << endl;

    return 0;
}

自分なりの解説

個人的に一番好きな問題です。ちょうど大学で線形代数とグラフ理論を少し学んでいる途中だったのですごくおもしろかったです。行列のn乗がこんなところで活用されてるということに驚きました。

行列aをijのパスがあれば1なければ0とする隣接行列を考える。この時a×aのij成分は
Σ $A_{i,k} \times A_{k,j}$
となり、iからkをとおってjにいくものをすべて足しているのでiからjにいく場合の数を表している。よって今回求めるのはパスがkのものを求めたいので行列をk乗する。ｋが非常に大きいので繰り返し二乗法を用いる。