sotanishy's competitive programming library

sotanishy's code snippets for competitive programming

View the Project on GitHub sotanishy/cp-library-cpp

:heavy_check_mark: Characteristic Polynomial
(math/linalg/characteristic_polynomial.hpp)

Description

正方行列の特性多項式を求める.

Operations

Reference

Depends on

Verified with

Code

#pragma once
#include <algorithm>
#include <vector>

#include "../polynomial.cpp"
#include "matrix.hpp"

template <typename mint>
Polynomial<mint> characteristic_polynomial(Matrix<mint> mat) {
    mat.assert_square();
    const int n = mat.row_size();
    if (n == 0) return {1};
    // stage 1: reduce mat to upper Hessenberg form
    for (int j = 0; j < n; ++j) {
        int i = j + 1;
        while (i < n && mat[i][j] == 0) ++i;
        if (i == n) continue;

        if (i != j + 1) {
            // swap mat[i], mat[j+1]
            mat[i].swap(mat[j + 1]);
            // swap mat[:,i], mat[:,j+1]
            for (int k = 0; k < n; ++k) {
                std::swap(mat[k][i], mat[k][j + 1]);
            }
        }

        mint inv = mint(1) / mat[j + 1][j];
        for (int k = j + 2; k < n; ++k) {
            mint v = mat[k][j] * inv;
            // mat[k] -= mat[j+1] * v
            for (int l = j; l < n; ++l) {
                mat[k][l] -= mat[j + 1][l] * v;
            }
            // mat[:,j+1] += mat[:,k] * v
            for (int l = 0; l < n; ++l) {
                mat[l][j + 1] += mat[l][k] * v;
            }
        }
    }
    // stage 2: recursively compute char polys of leading principal submatrices
    std::vector<Polynomial<mint>> p(n + 1);
    p[0] = {1};
    for (int i = 0; i < n; ++i) {
        p[i + 1].resize(i + 1 + 1);
        // (x - mat[i,i]) * p[i]
        for (int j = 0; j <= i; ++j) {
            p[i + 1][j] -= mat[i][i] * p[i][j];
            p[i + 1][j + 1] += p[i][j];
        }

        mint beta = 1;
        for (int k = i - 1; k >= 0; --k) {
            beta *= mat[k + 1][k];
            p[i + 1] -= p[k] * mat[k][i] * beta;
        }
    }
    return p[n];
}
#line 2 "math/linalg/characteristic_polynomial.hpp"
#include <algorithm>
#include <vector>

#line 3 "math/polynomial.cpp"
#include <cassert>
#line 5 "math/polynomial.cpp"

#line 2 "convolution/ntt.hpp"
#include <bit>
#line 4 "convolution/ntt.hpp"

constexpr int get_primitive_root(int mod) {
    if (mod == 167772161) return 3;
    if (mod == 469762049) return 3;
    if (mod == 754974721) return 11;
    if (mod == 998244353) return 3;
    if (mod == 1224736769) return 3;
}

template <typename mint>
void ntt(std::vector<mint>& a) {
    constexpr int mod = mint::mod();
    constexpr mint primitive_root = get_primitive_root(mod);

    const int n = a.size();
    for (int m = n; m > 1; m >>= 1) {
        mint omega = primitive_root.pow((mod - 1) / m);
        for (int s = 0; s < n / m; ++s) {
            mint w = 1;
            for (int i = 0; i < m / 2; ++i) {
                mint l = a[s * m + i];
                mint r = a[s * m + i + m / 2];
                a[s * m + i] = l + r;
                a[s * m + i + m / 2] = (l - r) * w;
                w *= omega;
            }
        }
    }
}

template <typename mint>
void intt(std::vector<mint>& a) {
    constexpr int mod = mint::mod();
    constexpr mint primitive_root = get_primitive_root(mod);

    const int n = a.size();
    for (int m = 2; m <= n; m <<= 1) {
        mint omega = primitive_root.pow((mod - 1) / m).inv();
        for (int s = 0; s < n / m; ++s) {
            mint w = 1;
            for (int i = 0; i < m / 2; ++i) {
                mint l = a[s * m + i];
                mint r = a[s * m + i + m / 2] * w;
                a[s * m + i] = l + r;
                a[s * m + i + m / 2] = l - r;
                w *= omega;
            }
        }
    }
}

template <typename mint>
std::vector<mint> convolution(std::vector<mint> a, std::vector<mint> b) {
    const int size = a.size() + b.size() - 1;
    const int n = std::bit_ceil((unsigned int)size);
    a.resize(n);
    b.resize(n);
    ntt(a);
    ntt(b);
    for (int i = 0; i < n; ++i) a[i] *= b[i];
    intt(a);
    a.resize(size);
    mint n_inv = mint(n).inv();
    for (int i = 0; i < size; ++i) a[i] *= n_inv;
    return a;
}
#line 7 "math/polynomial.cpp"

template <typename mint>
class Polynomial : public std::vector<mint> {
    using Poly = Polynomial;

   public:
    using std::vector<mint>::vector;
    using std::vector<mint>::operator=;

    Poly pre(int size) const {
        return Poly(this->begin(),
                    this->begin() + std::min((int)this->size(), size));
    }

    Poly rev(int deg = -1) const {
        auto ret = *this;
        if (deg != -1) ret.resize(deg, 0);
        return Poly(ret.rbegin(), ret.rend());
    }

    void trim() {
        while (!this->empty() && this->back() == 0) this->pop_back();
    }

    // --- unary operation ---

    Poly& operator-() const {
        auto ret = *this;
        for (auto& x : ret) x = -x;
        return ret;
    }

    // -- binary operation with scalar ---

    Poly& operator+=(const mint& rhs) {
        if (this->empty()) this->resize(1);
        (*this)[0] += rhs;
        return *this;
    }

    Poly& operator-=(const mint& rhs) {
        if (this->empty()) this->resize(1);
        (*this)[0] -= rhs;
        return *this;
    }

    Poly& operator*=(const mint& rhs) {
        for (auto& x : *this) x *= rhs;
        return *this;
    }

    Poly& operator/=(const mint& rhs) { return *this *= rhs.inv(); }

    Poly operator+(const mint& rhs) const { return Poly(*this) += rhs; }
    Poly operator-(const mint& rhs) const { return Poly(*this) -= rhs; }
    Poly operator*(const mint& rhs) const { return Poly(*this) *= rhs; }
    Poly operator/(const mint& rhs) const { return Poly(*this) /= rhs; }

    // --- binary operation with polynomial ---

    Poly& operator+=(const Poly& rhs) {
        if (this->size() < rhs.size()) this->resize(rhs.size());
        for (int i = 0; i < (int)rhs.size(); ++i) (*this)[i] += rhs[i];
        return *this;
    }

    Poly& operator-=(const Poly& rhs) {
        if (this->size() < rhs.size()) this->resize(rhs.size());
        for (int i = 0; i < (int)rhs.size(); ++i) (*this)[i] -= rhs[i];
        return *this;
    }

    Poly& operator*=(const Poly& rhs) {
        *this = convolution(*this, rhs);
        return *this;
        // // naive convolution O(N^2)
        // std::vector<mint> res(this->size() + rhs.size() - 1);
        // for (int i = 0; i < (int)this->size(); ++i) {
        //     for (int j = 0; j < (int)rhs.size(); ++j) {
        //         res[i + j] += (*this)[i] * rhs[j];
        //     }
        // }
        // return *this = res;
    }

    Poly& operator/=(const Poly& rhs) {
        if (this->size() < rhs.size()) {
            this->clear();
            return *this;
        }
        int n = this->size() - rhs.size() + 1;
        return *this = (rev().pre(n) * rhs.rev().inv(n)).pre(n).rev(n);
    }

    Poly& operator%=(const Poly& rhs) {
        *this -= *this / rhs * rhs;
        trim();
        return *this;
    }

    std::pair<Poly, Poly> divmod(const Poly& rhs) {
        auto q = *this / rhs;
        auto r = *this - q * rhs;
        r.trim();
        return {q, r};
    }

    Poly operator+(const Poly& rhs) const { return Poly(*this) += rhs; }
    Poly operator-(const Poly& rhs) const { return Poly(*this) -= rhs; }
    Poly operator*(const Poly& rhs) const { return Poly(*this) *= rhs; }
    Poly operator/(const Poly& rhs) const { return Poly(*this) /= rhs; }
    Poly operator%(const Poly& rhs) const { return Poly(*this) %= rhs; }

    // --- shift operation ---

    Poly operator<<(int n) const {
        auto ret = *this;
        ret.insert(ret.begin(), n, 0);
        return ret;
    }

    Poly operator>>(int n) const {
        if ((int)this->size() <= n) return {};
        auto ret = *this;
        ret.erase(ret.begin(), ret.begin() + n);
        return ret;
    }

    // --- evaluation ---

    mint operator()(const mint& x) {
        mint y = 0, powx = 1;
        for (int i = 0; i < (int)this->size(); ++i) {
            for (auto c : *this) {
                y += c * powx;
                powx *= x;
            }
            return y;
        }
    }

    // --- other operations ---

    Poly inv(int deg = -1) const {
        assert((*this)[0] != mint(0));
        if (deg == -1) deg = this->size();
        Poly res = {(*this)[0].inv()};
        for (int d = 1; d < deg; d <<= 1) {
            auto f = pre(2 * d);
            auto g = res;
            f.resize(2 * d);
            g.resize(2 * d);

            // g_{n+1} = g_n * (2 - g_n * f) mod x^{2^{n+1}}

            ntt(f);
            ntt(g);
            for (int i = 0; i < 2 * d; ++i) f[i] *= g[i];
            intt(f);

            for (int i = 0; i < d; ++i) f[i] = 0;

            ntt(f);
            for (int i = 0; i < 2 * d; ++i) f[i] *= g[i];
            intt(f);

            res.resize(2 * d);
            auto coef = mint(2 * d).inv().pow(2);
            for (int i = d; i < 2 * d; ++i) res[i] = -f[i] * coef;
        }
        return res.pre(deg);
    }

    Poly exp(int deg = -1) const {
        assert((*this)[0] == mint(0));
        if (deg == -1) deg = this->size();
        Poly ret = {mint(1)};
        for (int i = 1; i < deg; i <<= 1) {
            ret = (ret * (this->pre(i << 1) + mint(1) - ret.log(i << 1)))
                      .pre(i << 1);
        }
        return ret;
    }

    Poly log(int deg = -1) const {
        assert((*this)[0] == mint(1));
        if (deg == -1) deg = this->size();
        return (diff() * inv(deg)).pre(deg - 1).integral();
    }

    Poly pow(long long k, int deg = -1) const {
        if (k == 0) return {1};
        if (deg == -1) deg = this->size();
        auto ret = *this;
        int cnt0 = 0;
        while (cnt0 < (int)ret.size() && ret[cnt0] == 0) ++cnt0;
        if (cnt0 > (deg - 1) / k) return {};
        ret = ret >> cnt0;
        deg -= cnt0 * k;
        ret = ((ret / ret[0]).log(deg) * k).exp(deg) * ret[0].pow(k);
        ret = ret << (cnt0 * k);
        return ret;
    }

    Poly diff() const {
        Poly ret(std::max(0, (int)this->size() - 1));
        for (int i = 1; i <= (int)ret.size(); ++i)
            ret[i - 1] = (*this)[i] * mint(i);
        return ret;
    }

    Poly integral() const {
        Poly ret(this->size() + 1);
        ret[0] = mint(0);
        for (int i = 0; i < (int)ret.size() - 1; ++i)
            ret[i + 1] = (*this)[i] / mint(i + 1);
        return ret;
    }

    Poly taylor_shift(long long c) const {
        const int n = this->size();
        std::vector<mint> fact(n, 1), fact_inv(n, 1);
        for (int i = 1; i < n; ++i) fact[i] = fact[i - 1] * i;
        fact_inv[n - 1] = mint(1) / fact[n - 1];
        for (int i = n - 1; i > 0; --i) fact_inv[i - 1] = fact_inv[i] * i;

        auto ret = *this;
        Poly e(n + 1);
        e[0] = 1;
        mint p = c;
        for (int i = 1; i < n; ++i) {
            ret[i] *= fact[i];
            e[i] = p * fact_inv[i];
            p *= c;
        }
        ret = (ret.rev() * e).pre(n).rev();
        for (int i = n - 1; i >= 0; --i) {
            ret[i] *= fact_inv[i];
        }
        return ret;
    }
};
#line 4 "math/linalg/matrix.hpp"
#include <cmath>
#include <initializer_list>
#include <type_traits>
#line 8 "math/linalg/matrix.hpp"

template <typename T>
class Matrix {
   public:
    static Matrix concat(const Matrix& A, const Matrix& B) {
        assert(A.row == B.row);
        Matrix C(A.row, A.col + B.col);
        for (int i = 0; i < A.row; ++i) {
            std::ranges::copy(A[i], C[i].begin());
            std::ranges::copy(B[i], C[i].begin() + A.col);
        }
        return C;
    }

    static Matrix I(int n) {
        Matrix ret(n);
        for (int i = 0; i < n; ++i) ret[i][i] = 1;
        return ret;
    }

    Matrix() = default;
    Matrix(int n) : Matrix(n, n) {}
    Matrix(int row, int col)
        : row(row), col(col), mat(row, std::vector<T>(col)) {}
    Matrix(const std::vector<std::vector<T>>& mat)
        : row(mat.size()), col(mat[0].size()), mat(mat) {}

    int row_size() const { return row; }
    int col_size() const { return col; }
    std::pair<int, int> shape() const { return {row, col}; }

    const std::vector<T>& operator[](int i) const { return mat[i]; }
    std::vector<T>& operator[](int i) { return mat[i]; }

    Matrix submatrix(int i0, int i1, int j0, int j1) const {
        Matrix ret(i1 - i0, j1 - j0);
        for (int i = i0; i < i1; ++i) {
            std::ranges::copy(mat[i].begin() + j0, mat[i].begin() + j1,
                              ret.mat[i - i0].begin());
        }
        return ret;
    }

    // --- binary operations with matrix ---

    Matrix& operator+=(const Matrix& rhs) {
        assert(shape() == rhs.shape());
        for (int i = 0; i < row; ++i) {
            for (int j = 0; j < col; ++j) {
                mat[i][j] += rhs[i][j];
            }
        }
        return *this;
    }
    Matrix& operator-=(const Matrix& rhs) {
        assert(shape() == rhs.shape());
        for (int i = 0; i < row; ++i) {
            for (int j = 0; j < col; ++j) {
                mat[i][j] -= rhs[i][j];
            }
        }
        return *this;
    }
    Matrix& operator*=(const Matrix& rhs) {
        assert(col == rhs.row);
        Matrix res(row, rhs.col);
        for (int i = 0; i < row; ++i) {
            for (int k = 0; k < col; ++k) {
                for (int j = 0; j < rhs.col; ++j) {
                    res[i][j] += mat[i][k] * rhs.mat[k][j];
                }
            }
        }
        return *this = res;
    }

    Matrix operator+(const Matrix& rhs) const { return Matrix(*this) += rhs; }
    Matrix operator-(const Matrix& rhs) const { return Matrix(*this) -= rhs; }
    Matrix operator*(const Matrix& rhs) const { return Matrix(*this) *= rhs; }

    constexpr bool operator==(const Matrix& rhs) const {
        if (shape() != rhs.shape()) return false;
        for (int i = 0; i < row; ++i) {
            for (int j = 0; j < col; ++j) {
                if (!eq(mat[i][j], rhs.mat[i][j])) return false;
            }
        }
        return true;
    }

    // --- scalar multiplication ---

    Matrix& operator*=(const T& rhs) {
        for (auto& row : mat) {
            for (auto& x : row) x *= rhs;
        }
        return *this;
    }
    Matrix& operator/=(const T& rhs) { return *this /= rhs; }

    Matrix operator*(const T& rhs) const { return Matrix(*this) *= rhs; }
    Matrix operator/(const T& rhs) const { return Matrix(*this) /= rhs; }

    // --- other operations for general matrices ---

    Matrix operator-() const {
        Matrix ret(*this);
        for (auto& row : ret.mat) {
            for (auto& x : row) x = -x;
        }
        return ret;
    }

    Matrix transpose() const {
        Matrix ret(col, row);
        for (int i = 0; i < col; ++i) {
            for (int j = 0; j < row; ++j) {
                ret[i][j] = mat[j][i];
            }
        }
        return ret;
    }

    void reduce() {
        int pivot = 0;
        for (int j = 0; j < col; ++j) {
            int i = pivot;
            while (i < row && eq(mat[i][j], T(0))) ++i;
            if (i == row) continue;

            if (i != pivot) mat[i].swap(mat[pivot]);

            T pinv = T(1) / mat[pivot][j];
            for (int l = j; l < col; ++l) mat[pivot][l] *= pinv;

            for (int k = 0; k < row; ++k) {
                if (k == pivot) continue;
                T v = mat[k][j];
                for (int l = j; l < col; ++l) {
                    mat[k][l] -= mat[pivot][l] * v;
                }
            }

            ++pivot;
        }
    }

    int rank() const {
        auto A = *this;
        A.reduce();
        for (int i = 0; i < row; ++i) {
            bool nonzero = false;
            for (int j = 0; j < col; ++j) {
                if (!eq(A[i][j], T(0))) {
                    nonzero = true;
                    break;
                }
            }
            if (!nonzero) return i;
        }
        return row;
    }

    // --- other operations for square matrices ---

    void assert_square() const { assert(row == col); }

    Matrix pow(long long k) const {
        assert_square();
        auto ret = I(row);
        auto A = *this;
        while (k > 0) {
            if (k & 1) ret *= A;
            A *= A;
            k >>= 1;
        }
        return ret;
    }

    T det() const {
        assert_square();
        auto A = *this;
        T ret = 1;
        for (int j = 0; j < col; ++j) {
            int i = j;
            while (i < row && eq(A[i][j], T(0))) ++i;
            if (i == row) return 0;

            if (i != j) {
                A[i].swap(A[j]);
                ret = -ret;
            }

            T p = A[j][j];
            ret *= p;
            auto pinv = T(1) / p;
            for (int l = j; l < col; ++l) A[j][l] *= pinv;

            for (int k = j + 1; k < row; ++k) {
                T v = A[k][j];
                for (int l = j; l < col; ++l) {
                    A[k][l] -= A[j][l] * v;
                }
            }
        }
        return ret;
    }

    Matrix inv() const {
        assert_square();
        auto IB = concat(*this, I(row));
        IB.reduce();
        assert(IB.submatrix(0, row, 0, col) == I(row));
        return IB.submatrix(0, row, col, 2 * col);
    }

    template <typename U,
              typename std::enable_if<std::is_floating_point<U>::value>::type* =
                  nullptr>
    static constexpr bool eq(U a, U b) {
        return std::abs(a - b) < 1e-8;
    }

    template <typename U, typename std::enable_if<!std::is_floating_point<
                              U>::value>::type* = nullptr>
    static constexpr bool eq(U a, U b) {
        return a == b;
    }

   protected:
    int row, col;
    std::vector<std::vector<T>> mat;
};
#line 7 "math/linalg/characteristic_polynomial.hpp"

template <typename mint>
Polynomial<mint> characteristic_polynomial(Matrix<mint> mat) {
    mat.assert_square();
    const int n = mat.row_size();
    if (n == 0) return {1};
    // stage 1: reduce mat to upper Hessenberg form
    for (int j = 0; j < n; ++j) {
        int i = j + 1;
        while (i < n && mat[i][j] == 0) ++i;
        if (i == n) continue;

        if (i != j + 1) {
            // swap mat[i], mat[j+1]
            mat[i].swap(mat[j + 1]);
            // swap mat[:,i], mat[:,j+1]
            for (int k = 0; k < n; ++k) {
                std::swap(mat[k][i], mat[k][j + 1]);
            }
        }

        mint inv = mint(1) / mat[j + 1][j];
        for (int k = j + 2; k < n; ++k) {
            mint v = mat[k][j] * inv;
            // mat[k] -= mat[j+1] * v
            for (int l = j; l < n; ++l) {
                mat[k][l] -= mat[j + 1][l] * v;
            }
            // mat[:,j+1] += mat[:,k] * v
            for (int l = 0; l < n; ++l) {
                mat[l][j + 1] += mat[l][k] * v;
            }
        }
    }
    // stage 2: recursively compute char polys of leading principal submatrices
    std::vector<Polynomial<mint>> p(n + 1);
    p[0] = {1};
    for (int i = 0; i < n; ++i) {
        p[i + 1].resize(i + 1 + 1);
        // (x - mat[i,i]) * p[i]
        for (int j = 0; j <= i; ++j) {
            p[i + 1][j] -= mat[i][i] * p[i][j];
            p[i + 1][j + 1] += p[i][j];
        }

        mint beta = 1;
        for (int k = i - 1; k >= 0; --k) {
            beta *= mat[k + 1][k];
            p[i + 1] -= p[k] * mat[k][i] * beta;
        }
    }
    return p[n];
}
Back to top page