冥界の音楽 [yukicoder 1112]

https://yukicoder.me/problems/no/1112

前提知識

• 行列累乗によるDP更新最適化

解説

https://yukicoder.me/submissions/510596

知っていれば行列累乗が最初に候補に挙がってくる問題。
操作が10¹⁸回ある問題なので、logNにする必要はあり、数学的にパッと解ける問題にも見えないので、
とりあえず行列累乗から考え始めるといった感じ。
なお、本解説では行列累乗についてのちゃんとした解説はしないので、ググってどこかで勉強してきてほしい。

その前に

Nがもうちょっと少なかったらDPができる。
この1つ前の問題は似たような問題でDPで解くことができる。
dp[i][pre][lst] := i番目まで確定していて、1つ前がpreで、最後がlstである制約を満たしている曲の組合せ
iが10¹⁸でダメなのだが、遷移が毎回定数値を使っていて同じような感じになる場合は行列累乗による高速化が使える

行列累乗

dp[i][pre][lst] -> dp[i + 1][lst][nxt]の更新だが、行列累乗では1次元にする必要があるので、
dp[i][pre * K + lst] -> dp[i + 1][lst * K + nxt]のように先頭以外で1次元にしておこう。
すると、dp[pre * K + lst]に何か計算をして、dp[lst * K + nxt]を作ることになる。
計算前と計算後はベクトルであるので、ベクトルに対する計算、つまり行列計算を行う事で、1回の遷移が行える。
dp_next = M * dp_prev
このMは遷移可能であれば1, 遷移不可能なら0としておけばいい。
1回の遷移なら行列はMだし、2回の遷移なら行列はM²だし、N回の遷移なら行列はM^Nを使えばいい。
M^Nの行列計算は繰り返し二乗法を使えば効率的に計算可能。ここが『行列累乗』になる。
あとは、行列計算をうまいことやれば答え。

int K, M; ll N;
bool ok[6][6][6];
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> K >> M >> N;
    rep(i, 0, M) {
        int a, b, c;
        cin >> a >> b >> c;
        a--; b--; c--;
        ok[a][b][c] = true;
    }

    Vec v(K * K);
    rep(b, 0, K) rep(c, 0, K) if(ok[0][b][c]) v[b * K + c] = 1;

    Mat m(K * K, Vec(K * K));
    rep(a, 0, K) rep(b, 0, K) rep(c, 0, K) if (ok[a][b][c]) {
        m[b * K + c][a * K + b] = 1;
    }

    m = fastpow(m, N - 3);
    v = mulMatVec(m, v);

    int ans = 0;
    rep(a, 0, K) ans = add(ans, v[a * K + 0]);
    cout << ans << endl;
}