通过数组索引进行Numpy向量化循环

Question

我想做一个函数，将返回一个数组的booleans。基本上，代码将检查4-D数组中元素的相等性。它会遍历二维和三维并比较每个数组是否相等。下面是我当前的代码:

x = np.arange(0,72).reshape(4,2,3,3)
x[:,1,2,:] = 0
def equality(arr):
return np.array([np.allclose(arr[:,i,j,:],arr[:,i,j,:+1]) for i in range(arr.shape[1]) for j in range(arr.shape[2])])
print(x)
print(equality(x))

输出:

[[[[ 0  1  2]    [ 3  4  5]    [ 6  7  8]]

[[ 9 10 11]    [12 13 14]    [ 0  0  0]]]


[[[18 19 20]    [21 22 23]    [24 25 26]]

[[27 28 29]    [30 31 32]    [ 0  0  0]]]


[[[36 37 38]    [39 40 41]    [42 43 44]]

[[45 46 47]    [48 49 50]    [ 0  0  0]]]


[[[54 55 56]    [57 58 59]    [60 61 62]]

[[63 64 65]    [66 67 68]    [ 0  0  0]]]] [False False False False
False  True]

代码工作，但是有没有一种方法来摆脱使用numpy的东西，而不是for循环?

Answer 1

使用np.ndarray.reshape和numpy.transpose得到所需的形状。为了直观地了解顺序，检查这个问题。然后检查sub_arrays中的所有列是否相等:

>>> reshaped = np.transpose(x.reshape(4,6,3), axes=(1,0,2))
# 1st dim==4, 2nd*3rd dim==6, 4th dim==3, hence reshape(4,6,3)
>>> reshaped
array([[[ 0,  1,  2],
[18, 19, 20],
[36, 37, 38],
[54, 55, 56]],
[[ 3,  4,  5],
[21, 22, 23],
[39, 40, 41],
[57, 58, 59]],
[[ 6,  7,  8],
[24, 25, 26],
[42, 43, 44],
[60, 61, 62]],
[[ 9, 10, 11],
[27, 28, 29],
[45, 46, 47],
[63, 64, 65]],
[[12, 13, 14],
[30, 31, 32],
[48, 49, 50],
[66, 67, 68]],
[[ 0,  0,  0],
[ 0,  0,  0],
[ 0,  0,  0],
[ 0,  0,  0]]])
>>> (reshaped == reshaped[:, :, -1:]).all(1).all(1)
array([False, False, False, False, False,  True])