Question

假设我有一个3D数组（张量）

print a.shape
(100, 100, 100)

并希望使用一个数组索引其第一个轴：

print i.shape
(20,)

及其最后一个轴使用另一个数组

print j.shape
(30,)

我的意图是获得一个（20,100,30）数组，例如，用于

等作业

a[?!, :, ?!] = b

但是，我无法弄清楚如何。

print a[i.reshape(20, 1, 1), :, j.reshape(1, 1, 30)].shape
(20, 1, 30, 100)

print a[i.reshape(20, 1, 1), :, j].shape
(20, 1, 30, 100)

print a[i.reshape(20, 1), :, j].shape
(20, 30, 100)

print a[i.reshape(20, 1, 1), j.reshape(1, 1, 30)].shape
(20, 1, 30, 100)

据我了解＆＃34;高级索引＆＃34;规则，第一次尝试应该有效，但我甚至没有最终得到一个3D数组，并且完整维度（100）来自最终而不是中间。

Answer 1

方法＃1：您可以使用np.ix_来获取索引的这些网格，只需将索引编入输入数组就必须为您提供所需的输出。因此，实现就像这样 -

a[np.ix_(i,np.arange(a.shape[1]),j)]

方法＃2：使用广播索引的简单方法 -

a[i[:,None],:,j].swapaxes(1,2)

验证样本案例 -

In [24]: a = np.random.randint(0,99,(5,3,5))

In [25]: i = np.random.randint(0,5,(2))

In [26]: j = np.random.randint(0,5,(2))

In [27]: a[i[0],:,j[0]]
Out[27]: array([15,  7, 74])

In [28]: a[i[0],:,j[1]]
Out[28]: array([32, 19, 85])

In [29]: a[i[1],:,j[0]]
Out[29]: array([76, 65, 96])

In [30]: a[i[1],:,j[1]]
Out[30]: array([54, 65, 66])

In [31]: a[np.ix_(i,np.arange(a.shape[1]),j)]
Out[31]: 
array([[[15, 32],
        [ 7, 19],
        [74, 85]],

       [[76, 54],
        [65, 65],
        [96, 66]]])

In [50]: a[i[:,None],:,j].swapaxes(1,2)
Out[50]: 
array([[[15, 32],
        [ 7, 19],
        [74, 85]],

       [[76, 54],
        [65, 65],
        [96, 66]]])

使用索引编制值

对于方法＃1，它只是直截了当 -

a[np.ix_(i,np.arange(a.shape[1]),j)] = b

对于方法＃2，如果b是标量，它也应该是直截了当的 -

a[i[:,None],:,j] = b

对于方法＃2，如果您要分配形状b的ndarray (20,100,30)，我们需要在分配之前交换b的轴，如此 -

a[i[:,None],:,j] = np.swapaxes(b,1,2)

如何使用数组索引张量？

1 个答案: