Matrix multiplication and geometric sum (finite terms) code

[GaurangTandon]

Apologies for the spam. I intend to organize all this stuff sometime soon.

#include <bits/stdc++.h>
#ifdef ONLINE_JUDGE
#define endl "\n"
#endif
using namespace std;
typedef long long LL;
typedef vector<int> VI;
typedef vector<VI> VVI;
typedef pair<int, int> PII;

const int mod = 10;

VVI ide(int sz) {
    VVI res(sz, VI(sz, 0));
    for (int i = 0; i < sz; i++) res[i][i] = 1;
    return res;
}

VVI mult (const VVI &a, const VVI &b) {
    int r = a.size();
    int c = b.front().size();
    VVI res(r, VI(c, 0));
    int mid = a.front().size();
    
    for (int i = 0; i < r; i++) {
        for (int j = 0; j < c; j++) {
            for (int k = 0; k < mid; k++) {
                res[i][j] += 1ll * a[i][k] * b[k][j] % mod;
                res[i][j] %= mod;
            }
        }
    }

    return res;
}

VVI po(const VVI &mat, int y) {
    VVI res = ide(mat.size());
    VVI x = mat;

    while (y) {
        if (y & 1) {
            res = mult(res, x);
        }
        y >>= 1;
        x = mult(x, x);
    }

    return res;
}

// n is number of terms, gp starts with identity matrix
VVI gpsum(const VVI &mat, int n) {
    int r = mat.size(), c = mat.front().size();
    assert(r == c);

    if (n == 0) return VVI(r, VI(r, 0));
    if (n == 1) return ide(r);

    if (n & 1) {
        VVI mat2 = gpsum(mat, n - 1);
        mat2 = mult(mat, mat2);
        for (int i = 0; i < r; i++) mat2[i][i] += 1;
        return mat2;
    }

    VVI a = gpsum(mat, n / 2);
    VVI a2 = po(mat, n / 2);
    VVI b = mult(a2, a);
    for (int i = 0; i < r; i++)
        for (int j = 0; j < c; j++)
            a[i][j] += b[i][j], a[i][j] %= mod;
    return a;
}

void solve() {
    int n, k; cin >> n >> k;
    if (n == 0) exit(0);

    VVI mat(n, VI(n, 0));
    for (auto &x : mat) for (auto &y : x) cin >> y;

    mat = gpsum(mat, k + 1);
    for (int i = 0; i < n; i++) mat[i][i] -= 1, mat[i][i] += 10, mat[i][i] %= 10;

    for (auto &x : mat) { int i = -1; for (auto &y : x) ++i, cout << y << " \n"[i == n - 1]; }
    cout << endl;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
    while(true)
        solve();
    return 0;
}

Verified on UVa Power of a Matrix (11149).

[GaurangTandon]

Templated it, currently being run at kamil's gym contest (https://codeforces.com/gym/102644)

#include <bits/stdc++.h>
#ifdef ONLINE_JUDGE
#define endl "\n"
#endif
using namespace std;
typedef long long LL;
typedef vector<int> VI;
typedef vector<VI> VVI;
typedef pair<int, int> PII;

const int mod = 10;

template <typename T>
struct matrix {
    typedef vector<T> VT;
    typedef vector<VT> VVT;

    VVT ide(int sz) {
        VVT res(sz, VT(sz, 0));

        for (int i = 0; i < sz; i++) res[i][i] = 1;

        return res;
    }

    VVT mult (const VVT &a, const VVT &b) {
        int r = a.size();
        int c = b.front().size();
        VVT res(r, VT(c, 0));
        int mid = a.front().size();
        
        for (int i = 0; i < r; i++) {
            for (int j = 0; j < c; j++) {
                for (int k = 0; k < mid; k++) {
                    res[i][j] += a[i][k] * b[k][j];
                }
            }
        }

        return res;
    }

    VVT po(const VVT &mat, int y) {
        VVT res = ide(mat.size());
        VVT x = mat;

        while (y) {
            if (y & 1) {
                res = mult(res, x);
            }
            y >>= 1;
            x = mult(x, x);
        }

        return res;
    }
};

void solve() {
    long double p; int n; cin >> n >> p;

    vector<vector<long double>> term = { { 1 - p, p }, { p, 1 - p } };
    vector<vector<long double>> last = { { 1 }, { 0 } };

    matrix<long double> m;
    auto res = m.po(term, n);
    res = m.mult(res, last);

    cout << fixed << setprecision(10) << res[0][0] << endl;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
    solve();
    return 0;
}

[GaurangTandon]

template <typename T>
struct Matrix {
    typedef vector<T> VT;
    typedef vector<VT> VVT;
    VVT matrix;

    Matrix(const VVT &mat) {
        matrix = VVT(mat);
    }

	// identity matrix
    Matrix(int sz) {
        matrix = VVT(sz, VT(sz, 0));

        for (int i = 0; i < sz; i++) matrix[i][i] = 1;
    }

    Matrix(int r, int c) {
        matrix = VVT(r, VT(c, 0));
    }

    Matrix<T> mult(const Matrix<T> &bb) const {
        const auto &b = bb.matrix;
        const auto &a = matrix;

        int r = a.size();
        int c = b.front().size();
        Matrix<T> res(r, c);
        int mid = a.front().size();
        
        for (int i = 0; i < r; i++) {
            for (int j = 0; j < c; j++) {
                for (int k = 0; k < mid; k++) {
                    res.matrix[i][j] += 1ll * a[i][k] * b[k][j] % mod;
                    res.matrix[i][j] %= mod;
                }
            }
        }

        return res;
    }

	// in place multiplication
    void mult_(const VVT &b) {
        matrix = mult(b);
    }
};

Updated with latest style