时间戳

首页 > 解决方案 > 如何根据另一个二维数组中给出的索引对二维数组进行切片

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

问题描述

我有一个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不是循环遍历每一行？

标签： pythonnumpynumpy-ndarrayarray-broadcastingnumpy-slicing

解决方案

您可以利用数组索引并使用 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]

推荐阅读