Python 在数组的切片上迭代

假设我有一个3D numpy.array，例如，具有尺寸x y z，有没有一种方法可以沿着特定轴在切片上迭代？比如：

for layer in data.slices(dim=2):
    # do something with layer

编辑： 为了澄清，示例是一个dim=3数组，即shape=（len_x，len_y，len_z）。Elazar和相当于kamjagin的解决方案可以工作，但没有那么通用-您必须手动构造

[：，：，i]

，这意味着您需要知道维度，而代码不够通用，无法处理任意维度的数组。您可以使用类似于

[…，：]

的方法来填充缺少的维度，但是您仍然需要自己构建这个维度

对不起，应该更清楚一点，这个例子有点太简单了

类似的事

>>> data = np.array([[1,2,3],[4,5,6],[7,8,9]])
>>> for layer in [data[:,i] for i in range(3)]:
...     print layer
... 
[1 4 7]
[2 5 8]
[3 6 9]

---

如果我错了，请纠正我，但在我看来，您的3D阵列将看起来像：

>>> my_array.shape
    (3,N)

其中N是数组的大小。因此，如果您想在一个维度上进行迭代，只需执行以下操作：

>>> for item in my_array[1,:]:

---

这将在第二维度上进行迭代。

这可能比这更优雅地解决，但如果您事先知道dim（例如2），则有一种方法：

或者如果dim=0

for i in range(data.shape[dim]):
    layer = data[i,:,:]

---

等等。

在第一个维度上迭代非常容易，请参见下文。要在其他维度上迭代，请将该维度滚动到前面，并执行相同操作：

>>> data = np.arange(24).reshape(2, 3, 4)
>>> for dim_0_slice in data: # the first dimension is easy
...     print dim_0_slice
... 
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
[[12 13 14 15]
 [16 17 18 19]
 [20 21 22 23]]
>>> for dim_1_slice in np.rollaxis(data, 1): # for the others, roll it to the front
...     print dim_1_slice
... 
[[ 0  1  2  3]
 [12 13 14 15]]
[[ 4  5  6  7]
 [16 17 18 19]]
[[ 8  9 10 11]
 [20 21 22 23]]
>>> for dim_2_slice in np.rollaxis(data, 2):
...     print dim_2_slice
... 
[[ 0  4  8]
 [12 16 20]]
[[ 1  5  9]
 [13 17 21]]
[[ 2  6 10]
 [14 18 22]]
[[ 3  7 11]
 [15 19 23]]

---

编辑一些计时，以比较大型阵列的不同方法：

In [7]: a = np.arange(200*100*300).reshape(200, 100, 300)

In [8]: %timeit for j in xrange(100): a[:, j]
10000 loops, best of 3: 60.2 us per loop

In [9]: %timeit for j in xrange(100): a[:, j, :]
10000 loops, best of 3: 82.8 us per loop

In [10]: %timeit for j in np.rollaxis(a, 1): j
10000 loops, best of 3: 28.2 us per loop

In [11]: %timeit for j in np.swapaxes(a, 0, 1): j
10000 loops, best of 3: 26.7 us per loop

---

啊，对不起，应该更清楚一点-这是一个dim=3数组，即shape=（nx，ny，nz）+1，这就是想法，但您需要动态构建索引元组，例如

idx=（切片（无），）*dim+（i，）+（切片（无），）*（2-dim）

，然后

layer=data[idx]

。将维度滚动到一个设定的位置（第一个是最方便的），然后始终使用相同的迭代机制，这不太麻烦，请参见我的答案。啊，太好了。使用卷装解决方案绝对更优雅。请查看。swapaxes方法是最快的，但最不清晰。kamjagin的方法可以通过构造一个元组来推广（即，

数据[tuple（slice（None），slice（None），i）]

与

数据[：，：，i]

）是一样的。@AFoglia我认为这个问题的时间安排不是很相关。使用

swapaxes

或

rollaxis

将花费更多的时间来设置循环，但实际的迭代速度更快，请参见我的答案中添加的计时。在您非常小的示例中，设置控制着实际的迭代。我不同意可读性，但我可能太习惯numpy而没有注意到它。@AFoglia numpy有一个巧妙的技巧，可以将切片（元组）与要索引的对象分开构建：

numpy。s_[：，：，I]

相当于

元组（切片（无），切片（无），I）

。（请注意，

numpy.s

有一个属性，

maketuple

，该属性确定使用单个切片是否会产生元组[即

切片（开始、结束、步骤）

与

（切片（开始、结束、步骤））

。默认值为

False

）@JAB是的，但是如果他想要一个任意秩数组任意维的通用解决方案，那么用

（slice（None），*（idim）+（idx，）

编程生成一个元组要比用

：

字符更容易。你不能用

数据[（：，）*（idm）+（idx，）]

。哦，太好了，我不知道

rollaxis（）

。这里也一样，非常好！

In [7]: a = np.arange(200*100*300).reshape(200, 100, 300)

In [8]: %timeit for j in xrange(100): a[:, j]
10000 loops, best of 3: 60.2 us per loop

In [9]: %timeit for j in xrange(100): a[:, j, :]
10000 loops, best of 3: 82.8 us per loop

In [10]: %timeit for j in np.rollaxis(a, 1): j
10000 loops, best of 3: 28.2 us per loop

In [11]: %timeit for j in np.swapaxes(a, 0, 1): j
10000 loops, best of 3: 26.7 us per loop