nmWATI/Include/iAlg/iAlgMath/zxLinearSys.h

#pragma once

template <typename T>
static bool solve_linear_sys(T solution[], T* mat, int n, bool inplace);

// Gaussian elimination, n*(n+1) martix, O(n^3) time.
// m is array of ptr to each row
template <typename T>
static bool solve_linear_sys(T solution[], T* mat[], int n, bool inplace = false)
{
    if(!inplace)
    {
        int sz1 = sizeof(T) * (n+1);
        int sz = sz1 * n;
        T * mat1 = (T*) malloc(sz);
        T * p = mat1;
        for(int k=0; k<n; k++)
        {
            memcpy(p, mat[k], sz1);
            p+= n+1;
        }
        bool ret = solve_linear_sys(solution, mat1, n, true);
        free(mat1);
        return ret;
    }

    for(int k=0; k<n; k++)
    {
        // find pivot
        int r = k;
        T pivot = fabs(mat[k][k]);
        for(int i=k+1; i<n; i++)
        {
            T t = fabs(mat[i][k]);
            if(t > pivot)
            {
                pivot = t;
                r = i;
            }
        }

        if(pivot == 0)
        {
            return false;
        }

        if(r != k)
        {
            T* p = mat[r];
            mat[r] = mat[k];
            mat[k] = p;
        }

        // eliminate
        pivot = mat[k][k];
        for(int i=k+1; i<n; i++)
        {
            T a = mat[i][k];
            for(int j=k+1; j<=n; j++)
            {
                 mat[i][j] = mat[i][j] - mat[k][j] * a / pivot;
            }
            mat[i][k] = 0;
        }
    }

    // solution
    for(int k=n-1; k>=0; k--)
    {
        T c = 0;
        for(int j = k+1; j <n; j++)
        {
            c+= mat[k][j] * solution[j];
        }
        solution[k] = (mat[k][n] - c) / mat[k][k];
    }

    return true;
}

// m is n*(n+1) elements solid block
template <typename T>
static bool solve_linear_sys(T solution[], T* mat, int n, bool inplace = false)
{
    if(!inplace)
    {
        int sz = sizeof(T) * n * (n+1);
        T * mat1 = (T*) malloc(sz);
        memcpy(mat1, mat, sz);
        mat = mat1;
    }
    T ** m = (T**)malloc(sizeof(T*) * n);
    for(int i=0; i<n; i++)
    {
        m[i] = mat + i * (n+1);
    }
    bool ret = solve_linear_sys(solution, m, n, true);
    free(m);
    if(!inplace)
    {
        free(mat);
    }
    return ret;
}