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]