【C++】二维矩阵运算代码

C++ 二维矩阵运算代码

2023/09/04 by ShizuriYuki

该矩阵运算片段包含：

• 矩阵加法
• 矩阵减法
• 矩阵乘法
• 矩阵快速幂
• 矩阵求行列式
• 矩阵转置
• 矩阵求逆（高斯消元法）
• 矩阵求秩
• 矩阵拼接

//  Test of @matrix Snippet for C++
//  2023/09/04 by @ShizuriYuki

// #pragma GCC optimize (2)
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef unsigned long long ull;

// 矩阵加法
// mat_a: a * b, mat_b: a * b
template <typename T>
vector<vector<T>> matrixAdd(const vector<vector<T>>& _mat_a, const vector<vector<T>>& _mat_b) {
    if (_mat_a.size() != _mat_b.size() || _mat_a[0].size() != _mat_b[0].size()) {
        cerr << "[matrixAdd] Size not match: mat_a(" << _mat_a.size() << ", " << _mat_a[0].size()
             << "), mat_b(" << _mat_b.size() << ", " << _mat_b[0].size() << ")" << endl;
        throw exception();
    }
    size_t _a = _mat_a.size(), _b = _mat_a[0].size();
    vector<vector<T>> _mat_c(_a, vector<T>(_b, 0));
    for (size_t _idx_i = 0; _idx_i < _a; ++_idx_i) {
        for (size_t _idx_j = 0; _idx_j < _b; ++_idx_j) {
            _mat_c[_idx_i][_idx_j] = _mat_a[_idx_i][_idx_j] + _mat_b[_idx_i][_idx_j];
        }
    }
    return _mat_c;
}

// 矩阵减法
// mat_a: a * b, mat_b: a * b
template <typename T>
vector<vector<T>> matrixSub(const vector<vector<T>>& _mat_a, const vector<vector<T>>& _mat_b) {
    if (_mat_a.size() != _mat_b.size() || _mat_a[0].size() != _mat_b[0].size()) {
        cerr << "[matrixSub] Size not match: mat_a(" << _mat_a.size() << ", " << _mat_a[0].size()
             << "), mat_b(" << _mat_b.size() << ", " << _mat_b[0].size() << ")" << endl;
        throw exception();
    }
    size_t _a = _mat_a.size(), _b = _mat_a[0].size();
    vector<vector<T>> _mat_c(_a, vector<T>(_b, 0));
    for (size_t _idx_i = 0; _idx_i < _a; ++_idx_i) {
        for (size_t _idx_j = 0; _idx_j < _b; ++_idx_j) {
            _mat_c[_idx_i][_idx_j] = _mat_a[_idx_i][_idx_j] - _mat_b[_idx_i][_idx_j];
        }
    }
    return _mat_c;
}

// return: determinant of matrix
template <typename T>
long long matrixDet(const vector<vector<T>>& _mat_a) {
    size_t _a = _mat_a.size(), _b = _mat_a[0].size();
    if (_a != _b) {
        cerr << "[matrixDet] Size not match: (" << _a << ", " << _b << ")" << endl;
        throw exception();
    }
    long long _ans = 0;
    if (_a == 1) {
        _ans = _mat_a[0][0];
    }
    else {
        for (size_t _idx_i = 0; _idx_i < _a; ++_idx_i) {
            vector<vector<T>> _mat_b(_a - 1, vector<T>(_a - 1, 0));
            for (size_t _idx_j = 1; _idx_j < _a; ++_idx_j) {
                int _idx_l = 0;
                for (size_t _idx_k = 0; _idx_k < _a; ++_idx_k) {
                    if (_idx_k == _idx_i) {
                        continue;
                    }
                    _mat_b[_idx_j - 1][_idx_l++] = _mat_a[_idx_j][_idx_k];
                }
            }
            _ans += _mat_a[0][_idx_i] * matrixDet(_mat_b) * ((_idx_i & 1) ? -1 : 1);
        }
    }
    return _ans;
}

// return: rank of matrix
template <typename T>
int matrixRank(vector<vector<T>> _mat_a) {
    size_t _a = _mat_a.size(), _b = _mat_a[0].size();
    int _rank = 0;
    for (size_t _idx_i = 0; _idx_i < _b; ++_idx_i) {
        for (size_t _idx_j = _rank; _idx_j < _a; ++_idx_j) {
            if (_mat_a[_idx_j][_idx_i] != 0) {
                swap(_mat_a[_idx_j], _mat_a[_rank]);
                break;
            }
        }
        if (_mat_a[_rank][_idx_i] == 0) {
            continue;
        }
        for (size_t _idx_j = _rank + 1; _idx_j < _a; ++_idx_j) {
            if (_mat_a[_idx_j][_idx_i] == 0) {
                continue;
            }
            T _tmp = _mat_a[_idx_j][_idx_i] / _mat_a[_rank][_idx_i];
            for (size_t _idx_k = _idx_i; _idx_k < _b; ++_idx_k) {
                _mat_a[_idx_j][_idx_k] -= _mat_a[_rank][_idx_k] * _tmp;
            }
        }
        ++_rank;
    }
    return _rank;
}

// mat_a: a * b, mat_b: b * c
// _mod = 0: 不取模
// return: a * c
// template <typename T>
vector<vector<double>> matrixMul(const vector<vector<double>>& _mat_a, const vector<vector<double>>& _mat_b, double _mod = 0) {
    size_t _a = _mat_a.size(), _b = _mat_a[0].size(), _c = _mat_b[0].size();
    if (_b != _mat_b.size()) {
        cerr << "[matrixMul] Size not match: mat_a(" << _a << ", " << _b
             << "), mat_b(" << _mat_b.size() << ", " << _mat_b[0].size() << ")" << endl;
        throw exception();
    }
    vector<vector<double>> _mat_c(_a, vector<double>(_c, 0));
    for (size_t _idx_i = 0; _idx_i < _a; ++_idx_i) {
        for (size_t _idx_j = 0; _idx_j < _c; ++_idx_j) {
            for (size_t _idx_k = 0; _idx_k < _b; ++_idx_k) {
                _mat_c[_idx_i][_idx_j] += _mat_a[_idx_i][_idx_k] * _mat_b[_idx_k][_idx_j];
                // 考虑浮点数的情况，就不要用 %=_mod
                if (_mod) {
                    // _mat_c[_idx_i][_idx_j] %= _mod;
                }
            }
        }
    }
    return _mat_c;
}

// 矩阵快速幂
// 仅适用于整数，如需浮点数请自行修改
// _mod = 0: 不取模
// return: a ^ _hat
vector<vector<double>> matrixFastPow(const vector<vector<double>>& _mat_a, int _hat, int _mod = 0) {
    if (_mat_a.size() != _mat_a[0].size()) {
        cerr << "[matrixFastPow] Size not match: (" << _mat_a.size() << ", " << _mat_a[0].size() << ")" << endl;
        throw exception();
    }
    size_t _a = _mat_a.size();
    vector<vector<double>> _mat_c(_a, vector<double>(_a, 0));
    for (size_t _idx_i = 0; _idx_i < _a; ++_idx_i) {
        _mat_c[_idx_i][_idx_i] = 1;  // init as identity matrix, 1 on diagonal
    }
    vector<vector<double>> _mat_b = _mat_a;
    while (_hat) {
        if (_hat & 1) {
            _mat_c = matrixMul(_mat_c, _mat_b, _mod);
        }
        _mat_b = matrixMul(_mat_b, _mat_b, _mod);
        _hat >>= 1;
    }
    return _mat_c;
}

// 矩阵转置
// return: transpose matrix
template <typename T>
vector<vector<T>> matrixTranspose(const vector<vector<T>>& _mat_a) {
    size_t _a = _mat_a.size(), _b = _mat_a[0].size();
    vector<vector<T>> _mat_c(_b, vector<T>(_a, 0));
    for (size_t _idx_i = 0; _idx_i < _a; ++_idx_i) {
        for (size_t _idx_j = 0; _idx_j < _b; ++_idx_j) {
            _mat_c[_idx_j][_idx_i] = _mat_a[_idx_i][_idx_j];
        }
    }
    return _mat_c;
}

template <typename T>
vector<vector<T>> matrixConcat(const vector<vector<T>>& _mat_a, const vector<vector<T>> _mat_b, int axis = 0) {
    size_t _a = _mat_a.size(), _b = _mat_a[0].size();
    size_t _c = _mat_b.size(), _d = _mat_b[0].size();
    if (axis == 0) {
        if (_b != _d) {
            cerr << "[matrixConcat] Size not match: mat_a(" << _a << ", " << _b
                 << "), mat_b(" << _c << ", " << _d << ")" << endl;
            throw exception();
        }
        vector<vector<T>> _mat_c(_a + _c, vector<T>(_b));
        for (size_t _idx_i = 0; _idx_i < _a; ++_idx_i) {
            for (size_t _idx_j = 0; _idx_j < _b; ++_idx_j) {
                _mat_c[_idx_i][_idx_j] = _mat_a[_idx_i][_idx_j];
            }
        }
        for (size_t _idx_i = 0; _idx_i < _c; ++_idx_i) {
            for (size_t _idx_j = 0; _idx_j < _d; ++_idx_j) {
                _mat_c[_idx_i + _a][_idx_j] = _mat_b[_idx_i][_idx_j];
            }
        }
        return _mat_c;
    }
    else if (axis == 1) {
        if (_a != _c) {
            cerr << "[matrixConcat] Size not match: mat_a(" << _a << ", " << _b
                 << "), mat_b(" << _c << ", " << _d << ")" << endl;
            throw exception();
        }
        vector<vector<T>> _mat_c(_a, vector<T>(_b + _d));
        for (size_t _idx_i = 0; _idx_i < _a; ++_idx_i) {
            for (size_t _idx_j = 0; _idx_j < _b; ++_idx_j) {
                _mat_c[_idx_i][_idx_j] = _mat_a[_idx_i][_idx_j];
            }
        }
        for (size_t _idx_i = 0; _idx_i < _c; ++_idx_i) {
            for (size_t _idx_j = 0; _idx_j < _d; ++_idx_j) {
                _mat_c[_idx_i][_idx_j + _b] = _mat_b[_idx_i][_idx_j];
            }
        }
        return _mat_c;
    }
    else {
        cerr << "[matrixConcat] Invalid axis: " << axis << endl;
        throw exception();
    }
}

// 高斯消元法求逆矩阵
// 时间复杂度：O(n^3)；空间复杂度：O(n^2)
// 适用范围：矩阵的元素为float或double，且矩阵的行列数相等
template <typename T>
vector<vector<T>> matrixInverse_Gauss(vector<vector<T>> mat) {
    int n = mat.size();
    vector<vector<T>> inv = mat;
    vector<vector<T>> aug(n, vector<T>(n));
    for (int i = 0; i < n; i++)
        aug[i][i] = 1;
    for (int i = 0; i < n; i++) {
        // 找到第 i 列的最大元素所在的行
        int maxRow = i;
        for (int j = i + 1; j < n; j++) {
            if (abs(inv[j][i]) > abs(inv[maxRow][i])) {
                maxRow = j;
            }
        }
        // 交换第 i 行和第 maxRow 行
        for (int j = 0; j < n; j++) {
            swap(inv[i][j], inv[maxRow][j]);
            swap(aug[i][j], aug[maxRow][j]);
        }
        // 除以主元
        T temp = inv[i][i];
        for (int j = 0; j < n; j++) {
            inv[i][j] /= temp;
            aug[i][j] /= temp;
        }
        // 消元
        for (int j = 0; j < n; j++) {
            if (j != i) {
                temp = inv[j][i];
                for (int k = 0; k < n; k++) {
                    inv[j][k] -= inv[i][k] * temp;
                    aug[j][k] -= aug[i][k] * temp;
                }
            }
        }
    }
    return aug;
}

int main() {
    ios::sync_with_stdio(false);
    cout.tie(0);
    cin.tie(0);

    // vector<vector<double>> mat_a(3, vector<double>(3));
    // vector<vector<double>> mat_b(3, vector<double>(3));
    // // 输入矩阵
    // cout << "Input mat_a:" << endl;
    // for (int i = 0; i < 3; i++) {
    //     for (int j = 0; j < 3; ++j) {
    //         cin >> mat_a[i][j];
    //     }
    // }
    // cout << "Input mat_b:" << endl;
    // for (int i = 0; i < 3; i++) {
    //     for (int j = 0; j < 3; ++j) {
    //         cin >> mat_b[i][j];
    //     }
    // }

    vector<vector<double>> mat_a = {
        {2, 0, 0},
        {0, 2, 0},
        {0, 0, 2}};
    vector<vector<double>> mat_b = {
        {1, 1, 1},
        {0, 1, 1},
        {0, 0, 1}};

    // det(mat_a)
    cout << "det(mat_a): " << matrixDet(mat_a) << endl;
    // rank(mat_a)
    cout << "rank(mat_a): " << matrixRank(mat_a) << endl;
    // det(mat_b)
    cout << "det(mat_b): " << matrixDet(mat_b) << endl;
    // rank(mat_b)
    cout << "rank(mat_b): " << matrixRank(mat_b) << endl;

    // 求和
    vector<vector<double>> mat_c = matrixAdd(mat_a, mat_b);
    // 输出矩阵
    cout << "mat_a + mat_b -> mat_c:" << endl;
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; ++j) {
            cout << mat_c[i][j] << " ";
        }
        cout << endl;
    }

    // 求 c 的逆
    mat_c = matrixInverse_Gauss(mat_c);
    // 输出矩阵
    cout << "mat_c inverse -> mat_c:" << endl;
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; ++j) {
            cout << mat_c[i][j] << " ";
        }
        cout << endl;
    }

    // a * b
    mat_c = matrixMul(mat_a, mat_b);
    // 输出矩阵
    cout << "mat_a * mat_b -> mat_c:" << endl;
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; ++j) {
            cout << mat_c[i][j] << " ";
        }
        cout << endl;
    }

    // a ^ 3
    mat_c = matrixFastPow(mat_a, 3);
    // 输出矩阵
    cout << "mat_a ^ 3 -> mat_c:" << endl;
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; ++j) {
            cout << mat_c[i][j] << " ";
        }
        cout << endl;
    }

    // a 转置
    mat_c = matrixTranspose(mat_a);
    // 输出矩阵
    cout << "mat_a transpose -> mat_c:" << endl;
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; ++j) {
            cout << mat_c[i][j] << " ";
        }
        cout << endl;
    }

    // a 拼接 b 在 0 轴
    mat_c = matrixConcat(mat_a, mat_b, 0);
    // 输出矩阵
    cout << "mat_a concat mat_b at 0 -> mat_c:" << endl;
    for (int i = 0; i < 6; i++) {
        for (int j = 0; j < 3; ++j) {
            cout << mat_c[i][j] << " ";
        }
        cout << endl;
    }

    // a 拼接 b 在 1 轴
    mat_c = matrixConcat(mat_a, mat_b, 1);
    // 输出矩阵
    cout << "mat_a concat mat_b at 1 -> mat_c:" << endl;
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 6; ++j) {
            cout << mat_c[i][j] << " ";
        }
        cout << endl;
    }

    // c 转置
    mat_c = matrixTranspose(mat_c);
    // 输出矩阵
    cout << "mat_c transpose -> mat_c:" << endl;
    for (int i = 0; i < 6; i++) {
        for (int j = 0; j < 3; ++j) {
            cout << mat_c[i][j] << " ";
        }
        cout << endl;
    }

    return 0;
}

输出：

det(mat_a): 8
rank(mat_a): 3
det(mat_b): 1
rank(mat_b): 3
mat_a + mat_b -> mat_c:
3 1 1 
0 3 1 
0 0 3 
mat_c inverse -> mat_c:
0.333333 -0.111111 -0.0740741 
0 0.333333 -0.111111 
0 0 0.333333 
mat_a * mat_b -> mat_c:
2 2 2 
0 2 2 
0 0 2 
mat_a ^ 3 -> mat_c:
8 0 0 
0 8 0 
0 0 8 
mat_a transpose -> mat_c:
2 0 0 
0 2 0 
0 0 2 
mat_a concat mat_b at 0 -> mat_c:
2 0 0 
0 2 0 
0 0 2 
1 1 1 
0 1 1 
0 0 1 
mat_a concat mat_b at 1 -> mat_c:
2 0 0 1 1 1 
0 2 0 0 1 1 
0 0 2 0 0 1 
mat_c transpose -> mat_c:
2 0 0 
0 2 0 
0 0 2 
1 0 0 
1 1 0 
1 1 1