python - 如何根据另一个二维数组中给定的索引对二维数组进行切片

Question

我有一个 MxN名为 A 的数组存储我想要的数据。我还有一个 M x N2阵列 B存储数组索引，和 N2<N .每排B存储我想从该行的 A 获取的元素的索引。例如，以下代码对我有用:

A_reduced = np.zeros((M,N2))
for i in range(M):
    A_reduced[i,:] = A[i,B[i,:]]

是否有任何“矢量化”方法可以从 A 中提取所需元素基于 B而不是遍历每一行？

Answer 1

您可以利用数组索引并使用 reshape:

# set up M=N=4, N2=2
a = np.arange(16).reshape(4,4)
b = np.array([[1,2],[0,1],[2,3],[1,3]])

row_idx = np.repeat(np.arange(b.shape[0]),b.shape[1])
col_idx = b.ravel()

# output:
a[row_idx, col_idx].reshape(b.shape)

输出:

array([[ 1,  2],
       [ 4,  5],
       [10, 11],
       [13, 15]])

更新 : 另一个类似的解决方案

row_idx = np.repeat(np.arange(b.shape[0]),b.shape[1]).reshape(b.shape)

# output
a[row_idx,b]