Question

我有一个以numpy数组（A）作为输入的函数。根据数学计算，该阵列可以是2d或3d阵列。有一个整数m，可以是任何数字，除非数组是2D，m的值总是为0.我想将A的silce传递给另一个函数。由于A可以是3D或2D，我尝试了3D样式切片。

def fun(A):
    ... some code
    ans = fun2(A[:,:,m]) #The value of m is 0 if A is 2D

当A为2D

时，这会给我一个IndexError

IndexError: too many indices for array

如果A是2D，我想将完整的2D数组传递给fun2，就像在MATLAB中一样。如何在Python中完成？我使用Python 2。

Answer 1

似乎是一个很好的设置来使用np.atleast_3d因为我们可以强制它3D然后简单地沿着最后一个轴切割第m个索引，就像这样 -

np.atleast_3d(A)[...,m] # Or np.atleast_3d(A)[:,:,m]

它仍然是数组的视图，因此没有效率丢失！

案例运行

1）2D：

In [160]: A = np.random.randint(11,99,(4,5))

In [161]: np.atleast_3d(A)[...,0]
Out[161]: 
array([[13, 84, 38, 15, 26],
       [64, 91, 29, 11, 48],
       [25, 66, 77, 14, 87],
       [59, 96, 98, 30, 88]])

In [162]: A
Out[162]: 
array([[13, 84, 38, 15, 26],
       [64, 91, 29, 11, 48],
       [25, 66, 77, 14, 87],
       [59, 96, 98, 30, 88]])

2）3D：

In [163]: A = np.random.randint(11,99,(4,3,5))

In [164]: np.atleast_3d(A)[...,1]
Out[164]: 
array([[34, 81, 66],
       [56, 20, 25],
       [45, 36, 64],
       [82, 64, 31]])

In [165]: A[:,:,1]
Out[165]: 
array([[34, 81, 66],
       [56, 20, 25],
       [45, 36, 64],
       [82, 64, 31]])

在2D numpy数组中使用3D样式切片

1 个答案: