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

【问题标题】：Is there a way to conditionally index 3D-numpy array?有没有办法有条件地索引 3D-numpy 数组？
【发布时间】：2019-04-15 13:30:44
【问题描述】：

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

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

我要做的基本事情是计算 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)。

提前谢谢你，莱纳斯

【问题讨论】：

标签： python arrays numpy indexing

【解决方案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), :]

【讨论】：

【解决方案2】：

使用 np.take_along_axis 的替代解决方案（从版本 1.15 - docs）

import numpy as np
x = np.arange(2*6*6).reshape((2,6,6))
m = np.zeros(6, int)
m[0] = 1
#example: [1, 0, 0, 0, 0, 0]

np.take_along_axis(x, m[None, :, None], 0)   #add dimensions to mask to match array dimensions

>>array([[[36, 37, 38, 39, 40, 41],
        [ 6,  7,  8,  9, 10, 11],
        [12, 13, 14, 15, 16, 17],
        [18, 19, 20, 21, 22, 23],
        [24, 25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34, 35]]])

【讨论】：

谢谢，我会检查这个解决方案并与 blubberdiblub 的比较！