利用整数数组索引根据辅助数组索引提取元素

在给定的场景中，目标是从数组 A 使用第二个数组 B 中指定的索引。更直接的方法是使用 NumPy 的整数数组，而不是依赖 np.take 或 np.choose索引：

A[np.arange(A.shape[0]),B.ravel()]

以下是此代码如何实现所需结果：

• np.arange(A.shape[0]) 创建一个从 0 到 (A.shape[0] - 1) 的索引数组，表示A.
• B.ravel() 将 B 展平为一维数组，确保其元素数与 A 中的行数相同。
• 索引操作 A[... ] 使用 B 中的行索引和列索引检索 A 中的元素。

当 B 是一维数组或列索引列表时，此方法特别有用。通过跳过展平操作，代码变得更加简单：

A[np.arange(A.shape[0]),B]

示例：

A = np.array([[0, 1], [2, 3], [4, 5]])
B = np.array([1, 0, 1])

result = A[np.arange(A.shape[0]), B]

print(result)
# Output: [1, 2, 5]