能自动选取矩阵主元的代码

以下是一个能自动选取矩阵主元的 Python 代码示例：

import numpy as np

def gaussian_elimination(A, b):
    n = A.shape[0]
    Ab = np.hstack((A, b.reshape(n, 1)))

    for k in range(n):
        # 选取主元
        max_index = np.argmax(np.abs(Ab[k:, k])) + k
        Ab[[k, max_index]] = Ab[[max_index, k]]

        # 消元
        for i in range(k+1, n):
            factor = Ab[i, k] / Ab[k, k]
            Ab[i, k:] -= factor * Ab[k, k:]

    # 回带
    x = np.zeros(n)
    for i in range(n-1, -1, -1):
        x[i] = (Ab[i, -1] - Ab[i, :-1] @ x) / Ab[i, i]

    return x

在高斯消元的过程中，我们每次需要选取一个主元来进行消元。上述代码中的选主元步骤为：

max_index = np.argmax(np.abs(Ab[k:, k])) + k
Ab[[k, max_index]] = Ab[[max_index, k]]

其中，`np.argmax(np.abs(Ab[k:, k]))` 会返回从第 k 行开始，第 k 列及以后的所有元素的绝对值中的最大值的索引。由于 `np.argmax` 函数返回的是相对于切片的索引，因此需要加上 k 才能得到在整个矩阵中的索引。然后，我们将主元所在的行交换到第 k 行，以确保主元出现在对角线上。