如何在Python中实现线性方程组的回代法？修正代码输出错误

我正在Python中实现一个用于求解增广矩阵表示的线性方程组的回代法函数，已理解回代法概念，但实现细节出错。函数输入为3×3行阶梯形增广矩阵（右侧堆叠方程等号右边项），当前输出与期望结果不符。

原实现代码 #

# GRADED FUNCTION: back_substitution

def back_substitution(M):
    """
    Perform back substitution on an augmented matrix (with unique solution) in reduced row echelon form to find the solution to the linear system.

    Parameters:
    - M (numpy.array): The augmented matrix in row echelon form with unitary pivots (n x n+1).

    Returns:
    numpy.array: The solution vector of the linear system.
    """
    
    # Make a copy of the input matrix to avoid modifying the original
    M = M.copy()

    # Get the number of rows (and columns) in the matrix of coefficients
    num_rows = len(M)

    ### START CODE HERE ####
    
    # Iterate from bottom to top
    for row in range(num_rows): 
        # Get the substitution row. Remember now you must proceed from bottom to top, so you must start at the last row in the matrix.
        # The last row in matrix is in the index num_rows - 1, then you need to subtract the current index.
        substitution_row = M[row-1]

        # Iterate over the rows above the substitution_row
        for j in range(row + 1, num_rows): 

            # Get the row to be reduced. The indexing here is similar as above, with the row variable replaced by the j variable.
            row_to_reduce = M[j-1]
            # Get the index of the first non-zero element in the substitution row. This values does not depend on j!
            index = get_index_first_non_zero_value_from_row(substitution_row,row)

            # Get the value of the element at the found index
            value = row_to_reduce[index]

            # Perform the back substitution step using the formula row_to_reduce = None
            row_to_reduce = row_to_reduce - value*M[row]

            # Replace the updated row in the matrix, be careful with indexing!
            M[row] = row_to_reduce
            

    ### END CODE HERE ####

     # Extract the solution from the last column
    solution = M[:,-1]
    
    return solution

测试代码 #

A = np.array([[1,-1,0.5],[0,1,1], [0,0,1]])
B = np.array([[0.5], [-1], [-1]])
back_substitution(row_echelon_form(A,B))

当前输出 #

[-1,0,-1]

期望输出 #

[1,0,-1]

修正方案及解释 #

原代码的核心错误在于循环方向、索引逻辑以及行更新的目标位置完全错误，以下是修正后的代码：

# GRADED FUNCTION: back_substitution
import numpy as np

def get_index_first_non_zero_value_from_row(row, start_col):
    # 辅助函数：从start_col开始找行中第一个非零元素的索引
    for idx in range(start_col, len(row)-1):  # 最后一列是常数项，不参与主元查找
        if not np.isclose(row[idx], 0):
            return idx
    return -1

def back_substitution(M):
    """
    Perform back substitution on an augmented matrix (with unique solution) in reduced row echelon form to find the solution to the linear system.

    Parameters:
    - M (numpy.array): The augmented matrix in row echelon form with unitary pivots (n x n+1).

    Returns:
    numpy.array: The solution vector of the linear system.
    """
    
    # Make a copy of the input matrix to avoid modifying the original
    M = M.copy().astype(float)
    num_rows = M.shape[0]

    ### START CODE HERE ####
    # 回代法从最后一行（最下面的主元行）向上处理
    for row in range(num_rows - 1, -1, -1):
        # 获取当前主元行的主元索引
        pivot_idx = get_index_first_non_zero_value_from_row(M[row], row)
        # 处理当前主元行上方的所有行
        for upper_row in range(row - 1, -1, -1):
            # 获取上方行在主元位置的系数
            coeff = M[upper_row][pivot_idx]
            # 用主元行消去上方行的主元位置系数
            M[upper_row] -= coeff * M[row]
    ### END CODE HERE ####

    # 提取解向量
    solution = M[:, -1]
    return solution

关键修正点说明 #

1. 循环方向修正：原代码循环方向是从上到下，回代法需要从最后一行向上遍历到第0行，使用range(num_rows - 1, -1, -1)实现正确的遍历顺序。
2. 行更新目标修正：原代码错误地将更新后的行赋值给主元行，实际应该更新主元行上方的待消去行M[upper_row]。
3. 索引逻辑修正：原代码中substitution_row和row_to_reduce的索引完全混乱，修正后直接用当前主元行和上方行的索引操作矩阵，逻辑更清晰。
4. 辅助函数补全：补全get_index_first_non_zero_value_from_row的实现，确保能正确找到主元位置，同时用np.isclose处理浮点数精度问题。

用修正后的代码运行测试用例，将得到期望输出[1,0,-1]。

内容的提问来源于stack exchange，提问作者Zulfikar Irham