上三角线性方程组（线性方程组求解）

问题描述

在一个上三角线性方程组基础上，进行线性方程组求解。

输入形式

在屏幕上依次输入方阵阶数n，系数矩阵A和常数矩阵B。

输出形式

每一行输出一个根

样例1输入

4

4 -1 2 3

0 -2 7 -4

0 0 6 5

0 0 0 3

20

-7

4

6

样例1输出

[[ 3.]

[-4.]

[-1.]

[ 2.]]

样例1说明

输入：第1行为方阵阶数4，第2行至5行为系数矩阵A，第6行至9行为常数矩阵B。输出：每行依次输出方程解：x1, x2, x3, x4。

代码

# 回代法解上三角矩阵
import numpy as np

def backsub(A, B):
    n = B.size
    print(B.size)
    X = np.zeros([n,1], dtype = np.double)
    X[n-1] = B[n-1] / A[n-1, n-1]
    for k in range (n-2, -1, -1):
        X[k] = (B[k] - A[k, k+1:] @ X[k+1:]) / A[k, k]
    
    return X

def main():
    n = int(input())

    A = np.zeros([n,n],dtype=np.double)
    for r in range(n):
        A[r:] = np.array(input().split(),dtype = np.double)

    B = np.zeros((n,1),dtype=np.double)
    for r in range(n):
        B[r] = np.array(input(),dtype = np.double)
    
    X = backsub(A,B)
    print(X)

if __name__ == '__main__':
    main()