C++ – 矩阵与矩阵快速幂

struct mat {
    lli n, m;
    vector<lli> vec;

    mat(lli n, lli m) : n(n), m(m) { vec.resize(n * m); }

    lli& at(lli i, lli j)       { return vec[(i - 1) * m + j - 1]; }
    lli  at(lli i, lli j) const { return vec[(i - 1) * m + j - 1]; }

    void zero(lli x = 0) { fill(vec.begin(), vec.end(), x); }

    void unit(lli off = 0) {
        for (lli i = 1; i <= n; ++i)
            for (lli j = 1; j <= m; ++j)
                at(i, j) = i + off == j;
    }

    friend mat operator*(const mat& a, const mat& b) {
        mat res(a.n, b.m);
        assert(a.m == b.n);
        for (lli i = 1; i <= res.n; ++i)
            for (lli j = 1; j <= res.m; ++j)
                for (lli k = 1; k <= min(a.m, b.n); ++k)
                    res.at(i, j) += a.at(i, k) * b.at(k, j), res.at(i, j) %= mod;
        return res;
    }

    friend mat& operator*=(mat& a, const mat& b) {
        a = a * b;
        return a;
    }

    friend mat qpow(mat m, lli e) {
        mat res(m.n, m.m);
        res.unit();
        while (e) {
            if (e & 1) res *= m;
            m *= m;
            e >>= 1;
        }
        return res;
    }

    void print() const {
        for (lli i = 1; i <= n; ++i) {
            for (lli j = 1; j <= m; ++j)
                printf("%lld ", at(i, j));
            printf("\n");
        }
    }
};

发表回复