Add superfast matrix pow test

Author: adamant-pwn

cp-algo/algebra/poly.hpp

@@ -33,7 +33,7 @@ namespace cp_algo::algebra {
         poly_t mod_xk(size_t k) const {return poly::impl::mod_xk(*this, k);} // %= x^k
         poly_t mul_xk(size_t k) const {return poly::impl::mul_xk(*this, k);} // *= x^k
         poly_t div_xk(size_t k) const {return poly::impl::div_xk(*this, k);} // /= x^k
-        poly_t substr(size_t l, size_t r) const {return poly::impl::substr(*this, l, r);}
+        poly_t substr(size_t l, size_t k) const {return poly::impl::substr(*this, l, k);}
 
         poly_t operator *= (const poly_t &t) {fft::mul(a, t.a); normalize(); return *this;}
         poly_t operator * (const poly_t &t) const {return poly_t(*this) *= t;}

cp-algo/algebra/poly/impl/base.hpp

@@ -51,10 +51,10 @@ namespace cp_algo::algebra::poly::impl {
         }
         return std::vector(begin(p.a) + std::min<size_t>(k, p.a.size()), end(p.a));
     }
-    auto substr(auto const& p, size_t l, size_t r) {
+    auto substr(auto const& p, size_t l, size_t k) {
         return std::vector(
             begin(p.a) + std::min(l, p.a.size()),
-            begin(p.a) + std::min(r, p.a.size())
+            begin(p.a) + std::min(l + k, p.a.size())
         );
     }
     auto reverse(auto p, size_t n) {

cp-algo/linalg/frobenius.hpp

@@ -3,64 +3,74 @@
 #include "matrix.hpp"
 #include "../algebra/poly.hpp"
 #include <vector>
-namespace cp_algo::linalg {    
-    template<bool reduce = false>
-    auto frobenius_basis(auto const& A) {
+namespace cp_algo::linalg {
+    enum frobenius_mode {blocks, full};
+    template<frobenius_mode mode = blocks>
+    auto frobenius_form(auto const& A) {
         using matrix = std::decay_t<decltype(A)>;
         using base = matrix::base;
         using polyn = algebra::poly_t<base>;
         assert(A.n() == A.m());
         size_t n = A.n();
-        struct krylov {
-            std::vector<vec<base>> basis; 
-            polyn rec;
-        };
-        std::vector<krylov> blocks;
-        std::vector<vec<base>> reduced;
-        while(size(reduced) < n) {
+        std::vector<polyn> charps;
+        std::vector<vec<base>> basis, basis_init;
+        while(size(basis) < n) {
+            size_t start = size(basis);
             auto generate_block = [&](auto x) {
-                krylov block;
                 while(true) {
-                    vec<base> y = x | vec<base>::ei(n + 1, size(reduced));
-                    for(auto &it: reduced) {
+                    vec<base> y = x | vec<base>::ei(n + 1, size(basis));
+                    for(auto &it: basis) {
                         y.reduce_by(it);
                     }
                     y.normalize();
                     if(vec<base>(y[std::slice(0, n, 1)]) == vec<base>(n)) {
-                        block.rec = std::vector<base>(
-                            begin(y) + n + size(reduced) - size(block.basis),
-                            begin(y) + n + size(reduced) + 1
-                        );
-                        return std::pair{block, vec<base>(y[std::slice(n, n, 1)])};
+                        return polyn(std::vector<base>(begin(y) + n, end(y)));
                     } else {
-                        block.basis.push_back(x);
-                        reduced.push_back(y);
+                        basis_init.push_back(x);
+                        basis.push_back(y);
                         x = A.apply(x);
                     }
                 }
             };
-            auto [block, full_rec] = generate_block(vec<base>::random(n));
-            if constexpr (reduce) {
-                if(vec<base>(full_rec[std::slice(0, size(reduced), 1)]) != vec<base>(size(reduced))) {
-                    auto x = block.basis[0];
-                    size_t start = 0;
-                    for(auto &[basis, rec]: blocks) {
-                        polyn cur_rec = std::vector<base>(
-                            begin(full_rec) + start, begin(full_rec) + start + rec.deg()
-                        );
-                        auto shift = cur_rec / block.rec;
+            auto full_rec = generate_block(vec<base>::random(n));
+            // Extra trimming to make it block-diagonal (expensive)
+            if constexpr (mode == full) {
+                if(full_rec.mod_xk(start) != polyn()) {
+                    auto charp = full_rec.div_xk(start);
+                    auto x = basis_init[start];
+                    start = 0;
+                    for(auto &rec: charps) {
+                        polyn cur_rec = full_rec.substr(start, rec.deg());
+                        auto shift = cur_rec / charp;
                         for(int j = 0; j <= shift.deg(); j++) {
-                            x.add_scaled(basis[j], shift[j]);
+                            x.add_scaled(basis_init[start + j], shift[j]);
                         }
                         start += rec.deg();
                     }
-                    reduced.erase(begin(reduced) + start, end(reduced));
-                    tie(block, full_rec) = generate_block(x.normalize());
+                    basis.resize(start);
+                    basis_init.resize(start);
+                    full_rec = generate_block(x.normalize());
                 }
             }
-            blocks.push_back(block);
+            charps.push_back(full_rec.div_xk(start));
+        }
+        // Find transform matrices while we're at it...
+        if constexpr (mode == full) {
+            for(size_t i = 0; i < size(basis); i++) {
+                for(size_t j = i + 1; j < size(basis); j++) {
+                    basis[i].reduce_by(basis[j]);
+                }
+                basis[i].normalize();
+                basis[i] = vec<base>(
+                    basis[i][std::slice(n, n, 1)]
+                ) * (base(1) / basis[i][i]);
+            }
+            auto T = matrix::from_range(basis_init);
+            auto Tinv = matrix::from_range(basis);
+            return std::tuple{T, Tinv, charps};
+        } else {
+            return charps;
         }
-        return blocks;
     }
 };
-#endif // CP_ALGO_LINALG_FROBENIUS_HPP
+#endif // CP_ALGO_LINALG_FROBENIUS_HPP

cp-algo/linalg/matrix.hpp

@@ -42,6 +42,30 @@ namespace cp_algo::linalg {
             }
         }
 
+        static matrix block_diagonal(std::vector<matrix> const& blocks) {
+            size_t n = 0;
+            for(auto &it: blocks) {
+                assert(it.n() == it.m());
+                n += it.n();
+            }
+            matrix res(n);
+            n = 0;
+            for(auto &it: blocks) {
+                for(size_t i = 0; i < it.n(); i++) {
+                    res[n + i][std::slice(n, it.n(), 1)] = it[i];
+                }
+                n += it.n();
+            }
+            return res;
+        }
+        static matrix random(size_t n, size_t m) {
+            matrix res(n, m);
+            std::ranges::generate(res, std::bind(vec<base>::random, m));
+            return res;
+        }
+        static matrix random(size_t n) {
+            return random(n, n);
+        }
         static matrix eye(size_t n) {
             matrix res(n);
             for(size_t i = 0; i < n; i++) {
@@ -97,24 +121,15 @@ namespace cp_algo::linalg {
             return bpow(*this, k, eye(n()));
         }
 
-        static matrix random(size_t n, size_t m) {
-            matrix res(n, m);
-            std::ranges::generate(res, std::bind(vec<base>::random, m));
-            return res;
-        }
-        static matrix random(size_t n) {
-            return random(n, n);
-        }
-
         matrix& normalize() {
             for(auto &it: *this) {
                 it.normalize();
             }
             return *this;
         }
 
-        enum Mode {normal, reverse};
-        template<Mode mode = normal>
+        enum gauss_mode {normal, reverse};
+        template<gauss_mode mode = normal>
         matrix& gauss() {
             for(size_t i = 0; i < n(); i++) {
                 row(i).normalize();
@@ -126,11 +141,11 @@ namespace cp_algo::linalg {
             }
             return normalize();
         }
-        template<Mode mode = normal>
+        template<gauss_mode mode = normal>
         auto echelonize(size_t lim) {
             return gauss<mode>().sort_classify(lim);
         }
-        template<Mode mode = normal>
+        template<gauss_mode mode = normal>
         auto echelonize() {
             return echelonize<mode>(m());
         }

cp-algo/linalg/vector.hpp

@@ -24,8 +24,6 @@ namespace cp_algo::linalg {
         bool operator != (vec const& t) const {return !(*this == t);}
 
         vec operator-() const {return Base::operator-();}
-        vec operator-(vec const& t) const {return Base::operator-(t);}
-        vec operator+(vec const& t) const {return Base::operator+(t);}
 
         static vec from_range(auto const& R) {
             vec res(std::ranges::distance(R));
@@ -45,9 +43,6 @@ namespace cp_algo::linalg {
             return res;
         }
 
-        // Make sure the result is vec, not Base
-        vec operator*(base t) const {return Base::operator*(t);}
-
         void add_scaled(vec const& b, base scale, size_t i = 0) {
             assert(false);
             for(; i < size(*this); i++) {

verify/algebra/matrix/characteristic.test.cpp

@@ -18,12 +18,8 @@ void solve() {
     cin >> n;
     matrix<base> A(n);
     A.read();
-    auto blocks = frobenius_basis(A);
-    polyn res(1);
-    for(auto &[basis, rec]: blocks) {
-        res *= rec;
-    }
-    res.print();
+    auto blocks = frobenius_form(A);
+    reduce(begin(blocks), end(blocks), polyn(1), multiplies{}).print();
 }
 
 signed main() {
@@ -35,4 +31,4 @@ signed main() {
     while(t--) {
         solve();
     }
-}
+}

verify/algebra/matrix/pow_fast.test.cpp

@@ -0,0 +1,52 @@
+// @brief Pow of Matrix (Frobenius)
+#define PROBLEM "https://judge.yosupo.jp/problem/pow_of_matrix"
+#pragma GCC optimize("Ofast,unroll-loops")
+#pragma GCC target("avx2,tune=native")
+#include "cp-algo/linalg/frobenius.hpp"
+#include <bits/stdc++.h>
+
+using namespace std;
+using namespace cp_algo::algebra;
+using namespace cp_algo::linalg;
+
+const int mod = 998244353;
+using base = modular<mod>;
+using polyn = poly_t<base>;
+
+template<typename base>
+auto frobenius_pow(matrix<base> A, uint64_t k) {
+    using polyn = poly_t<base>;
+    auto [T, Tinv, charps] = frobenius_form<full>(A);
+    vector<matrix<base>> blocks;
+    for(auto charp: charps) {
+        matrix<base> block(charp.deg());
+        auto xk = polyn::xk(1).powmod(k, charp);
+        for(size_t i = 0; i < block.n(); i++) {
+            ranges::copy(xk.a, begin(block[i]));
+            xk = xk.mul_xk(1) % charp;
+        }
+        blocks.push_back(block);
+    }
+    auto S = matrix<base>::block_diagonal(blocks);
+    return Tinv * S * T;
+}
+
+void solve() {
+    size_t n;
+    uint64_t k;
+    cin >> n >> k;
+    matrix<base> A(n);
+    A.read();
+    frobenius_pow(A, k).print();
+}
+
+signed main() {
+    //freopen("input.txt", "r", stdin);
+    ios::sync_with_stdio(0);
+    cin.tie(0);
+    int t = 1;
+    //cin >> t;
+    while(t--) {
+        solve();
+    }
+}