python - 有没有办法有条件地索引 3D-numpy 数组？

Question

有一个形状为 (2,6, 60) 的数组 A，是否可以根据形状为 (6,) 的二进制数组 B 对其进行索引？

6 和 60 非常随意，它们只是我希望访问的 2D 数据。

我试图做的基本事情是计算二维数据的两个变体(在本例中为 (6,60))，然后有效地选择总和最低的那些 -这就是二进制 (6,) 数组的来源。

示例:对于 B = [1,0,1,0,1,0] 我希望接收的内容等于堆叠

A[1,0,:]
A[0,1,:]
A[1,2,:]
A[0,3,:]
A[1,4,:]
A[0,5,:]

但我想通过直接索引而不是 for 循环来实现。

我已经尝试过 A[B]、A[:,B,:]、A[B,:,:] A[:,:,B] 但没有一个提供所需的(6,60) 矩阵。

import numpy as np
A = np.array([[4, 4, 4, 4, 4, 4], [1, 1, 1, 1, 1, 1]])
A = np.atleast_3d(A)
A = np.tile(A, (1,1,60)
B = np.array([1, 0, 1, 0, 1, 0])
A[B]

预期结果是一个 (6,60) 数组，其中包含上述 A 中的元素，接收到的是 (2,6,60) 或 (6,6,60)。

先谢谢你，莱纳斯

Answer 1

您可以生成一系列要迭代的索引，在您的例子中是从 0 到 5:

count = A.shape[1]

indices = np.arange(count)  # np.arange(6) for your particular case

>>> print(indices)
array([0, 1, 2, 3, 4, 5])

然后您可以使用它来进行高级索引:

result_array = A[B[indices], indices, :]

如果您始终按递增顺序使用 A 第二轴的从 0 到长度 - 1(即在您的情况下为 0 到 5)的完整范围，您可以将其简化为:

result_array = A[B, indices, :]
# or the ugly result_array = A[B, np.arange(A.shape[1]), :]

如果总是 6，甚至这样:

result_array = A[B, np.arange(6), :]