Matrix Computation

#include<bits/stdc++.h>
#define ll long long

// Support:
//   Matrix multiplication
//   Solving system of equations
template <typename T> struct Matrix 
{ 
    vector< vector<T> > data; 
    int nrow, ncol;
    T MMOD;

    Matrix (int nrow_, int ncol_, T MMOD_=-1) : nrow(nrow_), ncol(ncol_), MMOD(MMOD_) {

        data.clear();
        data.resize(nrow, vector<T>(0, 0));
        for (int i = 0; i < nrow; ++i) {
            data[i].resize(ncol, 0);
        }
    }
    
    Matrix (int nrow_, int ncol_, vector< vector<T> > data_, T MMOD_=-1) : nrow(nrow_), ncol(ncol_), data(data_), MMOD(MMOD_) {}
    
    Matrix<T> (const Matrix<T> &mt) {
        nrow = mt.nrow;
        ncol = mt.ncol;
        data = mt.data;
        MMOD = mt.MMOD;
    }
    
    static Matrix identityMatrix(int nrow, int ncol, T MMOD_) {
        Matrix<T> m (nrow, ncol, MMOD_);
        for (int i = 0; i < nrow; ++i)
            m.data[i][i] = 1;
        return m;
    }
    
    Matrix operator * (Matrix b) const {
        
        if (ncol != b.nrow) {
            cout << "multiplication error: matrices' size not valid \n";
            return Matrix(0, 0, MMOD);
        }
      
        Matrix<T> res(nrow, b.ncol, MMOD);
        for (int i = 0; i < nrow; ++i) 
            for (int j = 0; j < b.ncol; ++j)
                for (int k = 0; k < ncol; ++k) {
                    
                    res.data[i][j] += data[i][k]*b.data[k][j];
                
                    if (MMOD > 0 && (res.data[i][j] > MMOD || res.data[i][j] < 0))
                        res.data[i][j] = ((res.data[i][j] % MMOD) + MMOD) % MMOD;
                }
        
        return res;
    }
    
    Matrix pow(ll exp) {
        
        if (ncol != nrow) {
            cout << "Power error: matrices' size not valid \n";
            return Matrix(0, 0, MMOD);
        }
        
        if (exp == 0) 
            return identityMatrix(nrow, ncol, MMOD);
        
        if (exp == 1)
//          return (*this);
            return Matrix<T> (nrow, ncol, data, MMOD);
            
        Matrix res = pow(exp >> 1);
        res = res * res;
        if (exp & 1)
            res = res * (*this);
//      cout << "exp: " << exp << endl;
//      res.print();
        return res;
    }
    
    vector<double> solveSystemEquations() {
        // Use Gauss-Elimination
        // Intend to solve system of linear equations of form:
        //
        //      a[0][0]*x[0]   + a[0][1]*x[1]   + a[0][2]*x[2]   + ... + a[0][n-1]*x[n-1]   = a[0][n]
        //      a[1][0]*x[0]   + a[1][1]*x[1]   + a[1][2]*x[2]   + ... + a[1][n-1]*x[n-1]   = a[1][n]
        //                                              ...
        //      a[n-1][0]*x[0] + a[n-1][1]*x[1] + a[n-1][2]*x[2] + ... + a[n-1][n-1]*x[n-1] = a[n-1][n]
        //
        // (n rows and n+1 columns)
        
        // pre-check
        T tmp = 1.1;
        if (tmp == 1) {
            cout << "Matrix is of type int (or long long) which is not suitable for gauss-elimination \n";
        }
        
        auto gdata = data;
                
        // do gauss-elimination
        for (int i = 0; i < nrow; ++i) {
            int pivot = i;
            for (int j = i+1; j < nrow; ++j)
                if (fabs(gdata[j][i]) > fabs(gdata[pivot][i]))
                    pivot = j;
            
            if (pivot != i)
                swap(gdata[i], gdata[pivot]);
            
            if (fabs(gdata[i][i]) < EPS)
                continue;
            
            for (int k = ncol-1; k >= i; k--)
                gdata[i][k] /= gdata[i][i];
            
            for (int j = i+1; j < nrow; ++j) {
                for (int k = ncol-1; k >= i; --k)
                    gdata[j][k] -= gdata[j][i] * gdata[i][k];
            }
        }
        
        // find x[0..n-1]
        for (int i = nrow-2; i >= 0; --i) {
            if (fabs(gdata[i][i]) > EPS)  // means: if gdata[i][i] == 1
                for (int k = i+1; k < ncol-1; ++k) 
                    if (isfinite(gdata[i][k]))
                        gdata[i][ncol-1] -= gdata[i][k] * gdata[k][ncol-1];
            else {
                gdata[i][ncol-1] = NAN;
            }
        }
        
        vector<double> x;
        for (int i = 0; i < nrow; ++i) {
            x.push_back(gdata[i][ncol-1]);
        }

        return x;
        
    }
    
    void print() {
        cout << "print matrix " << nrow << " x " << ncol << endl;
        for (int i = 0; i < nrow; ++i) {
            for (int j = 0; j < ncol; ++j)
                cout << data[i][j] << " ";
            cout << "\n";
        }
    }
  
};

Leave a Reply