给我一段列主消元法求解的python代码

以下是使用列主消元法求解线性方程组的Python代码：

import numpy as np

def gaussian_elimination(A, b):
    n = len(A)
    # 构造增广矩阵
    Ab = np.concatenate((A, b.reshape(n, 1)), axis=1)
    # 高斯消元
    for i in range(n):
        # 找到主元所在的行
        max_row = i
        for j in range(i + 1, n):
            if abs(Ab[j, i]) > abs(Ab[max_row, i]):
                max_row = j
        # 将主元所在行交换到当前行
        Ab[[i, max_row]] = Ab[[max_row, i]]
        # 将当前行的主元归一化
        Ab[i] = Ab[i] / Ab[i, i]
        # 用当前行消元下面的所有行
        for j in range(i + 1, n):
            Ab[j] = Ab[j] - Ab[i] * Ab[j, i]
    # 回带求解
    x = np.zeros(n)
    for i in range(n - 1, -1, -1):
        x[i] = Ab[i, n] - np.dot(Ab[i, i+1:n], x[i+1:n])
    return x

其中，`A` 是 $n \times n$ 的系数矩阵，`b` 是 $n \times 1$ 的常数向量，返回值是 $n \times 1$ 的未知量向量 `x`。这个函数实现了高斯消元的过程，包括选主元、归一化和消元。在高斯消元完成后，再通过回带法求解出未知量向量 `x`。