'''
Author: SJ2050
Date: 2021-10-29 18:46:40
LastEditTime: 2021-10-29 22:41:31
Version: v0.0.1
Description: Solution for homework1.4.
Copyright © 2021 SJ2050
'''
import numpy as np
from scipy.linalg import lu_factor, lu_solve

def gauss_eliminate(K, b):
    K = np.array(K)
    b = np.array(b)
    n = len(K)
    x = np.zeros(n)

    for i in range(0, n-1):
        for j in range(i+1,n):
            c = -K[j][i]/K[i][i]
            for k in range(0, n):
                K[j][k] += K[i][k]*c
            b[j] += b[i]*c
    x[n-1]=b[n-1]/K[n-1][n-1]

    for i in range(n-2, -1, -1):
        for j in range(i+1, n):
            b[i] -= K[i][j]*x[j]
        x[i] = b[i]/K[i][i]

    return x


if __name__ == '__main__':
    K = np.mat([[3, 4, 2],
                [5, 3, 4],
                [8, 2, 7]])
    b = np.mat([[10], [14], [20]])
    # 1. 系数矩阵求逆求解线性方程组
    s1 = K**(-1)*b
    print('1.系数矩阵求逆求解：\n', s1)

    # 2. 直接求解法求解线性方程组
    s2 = np.linalg.solve(K, b)
    print('2.直接求解法解线性方程组：\n', s2)

    # 3. 使用LU分解求解线性方程组
    lu, piv = lu_factor(K)
    s3 = lu_solve((lu, piv), b)
    print('3.LU分解求解线性方程组: \n', s3)

    # 4. 使用高斯消元法求解线性方程组
    s4 = gauss_eliminate(K[:, :]*1.0, [x[0]*1.0 for x in b])
    print('4.高斯消元法求解线性方程组: \n', s4)