python - 如何从其他数组的不同起点对 2-D Numpy 数组的范围进行索引

Question

这个问题已经有答案了:

Selecting Random Windows from Multidimensional Numpy Array Rows (2 个回答)

已关闭 4 年前。

我正在尝试提取形状为 (3,32) 的 numpy 数组的子集并对其执行数学运算，我尝试提取的数据子集是形状 (3,9) 的数据子集，并且该数据的范围源自另一个数组大小 (3) 中包含的索引。例如，我有一个来自在时域中运行的三个 channel 的值的数据集，并将每个 channel 的最大值的索引提取到一个数组中

a = np.random.randint(20,size = (3,32))
a
array([[18,  3, 10,  6, 12,  1, 10,  8,  4, 11, 13, 14,  9,  9, 10,  2,
         9,  0,  0, 16, 14, 19,  1, 19, 14, 19, 19,  2, 14,  0,  4, 18],
       [ 9, 19,  2, 12,  0, 14, 18,  7,  3,  0,  7,  3, 12, 19,  4,  2,
         5,  9,  2, 11, 15, 19, 16, 17,  3,  4, 17,  5,  6,  1,  2, 17],
       [ 0, 11, 18,  8,  9,  2,  9, 15,  9,  6,  0,  8,  9, 16,  9,  6,
         1, 19,  1,  9, 12,  8,  0,  0,  7, 15,  3, 14, 15,  8, 10, 19]])
b = np.argmax(a,1)
b
array([21,  1, 17], dtype=int64)

我此时的目标是派生一个新数组，该数组由每个指定索引的三个值组成。例如，我希望提取出:

[21,22,23] from a[0]
[1,2,3] from a[1]
[17,18,19] from a[2]

全部放入大小为[3,3]的新数组

我已经能够使用循环来实现这一点，但我怀疑有一种更有效的方法可以在没有循环的情况下产生这个结果(速度是这个应用程序的一个问题)。我已经能够通过手动填充较小的矩阵来实现类似的结果......

c = np.asarray([[1,2,3],[2,3,4],[3,4,5]])
a[np.arange(3)[:,None],c]
array([[ 3, 10,  6],
       [ 2, 12,  0],
       [ 8,  9,  2]])

但是，考虑到此应用程序的动态性质，我想编写此应用程序，以便可以动态扩展它(索引范围为超出根索引的 9 个值等)。我只是不知道是否有这样的方法可以做到这一点。我使用了类似于以下的语法来对数组进行切片...

a[np.arange(3)[:,None],b[:]:(b[:] + 2)]

导致错误消息的性质......

builtins.TypeError: only integer scalar arrays can be converted to a scalar index

Answer 1

既然你说你不能溢出，那么这就变得不那么棘手了。一般来说，由于您有起始索引，因此您可以使用基本广播来使用索引创建 (n, 3) 形状数组，并使用 take_along_axis 来提取这些元素来自原始数组。

<小时/>

np.take_along_axis(a, b[:, None] + np.arange(3), axis=1)

array([[19,  1, 19],
       [19,  2, 12],
       [19,  1,  9]])