Python写列主元素高斯消去法

以下是Python实现列主元素高斯消去法的代码：

import numpy as np

# 列主元素高斯消去法
def gauss_elimination(A, b):
    n = len(b)
    for i in range(n):
        # 首先选取第i列中绝对值最大的元素作为主元素
        max_idx = i
        for j in range(i+1, n):
            if abs(A[j, i]) > abs(A[max_idx, i]):
                max_idx = j
        # 交换第i行和第max_idx行
        A[[i, max_idx], :] = A[[max_idx, i], :]
        b[[i, max_idx]] = b[[max_idx, i]]
        # 消元过程
        for j in range(i+1, n):
            factor = A[j, i] / A[i, i]
            A[j, i:] -= factor * A[i, i:]
            b[j] -= factor * b[i]
    # 回代求解
    x = np.zeros(n)
    for i in range(n-1, -1, -1):
        x[i] = (b[i] - np.dot(A[i, i+1:], x[i+1:])) / A[i, i]
    return x

以上代码中，`A`表示系数矩阵，`b`表示常数向量。函数中的注释已经解释了主要的步骤。