遍历多维数组的所有一维子数组

Question

在Python中，如何最快地遍历一个n维数组中的所有一维子数组呢？

比如说，考虑一个三维数组：

import numpy as np 
a = np.arange(24)
a = a.reshape(2,3,4)

我们希望从这个迭代器中得到的结果顺序是：

a[:,0,0]
a[:,0,1]
..
a[:,2,3]
a[0,:,0]
..
a[1,:,3]
a[0,0,:]
..
a[1,2,:]

Answer 1

你的朋友们是 slice() 对象、numpy 的 ndarray.__getitem__() 方法，还有可能是 itertools.chain.from_iterable。

Answer 2

可能还有更有效的方法，但这个方法应该可以用...

import itertools
import numpy as np

a = np.arange(24)
a = a.reshape(2,3,4)

colon = slice(None)
dimensions = [range(dim) + [colon] for dim in a.shape]

for dim in itertools.product(*dimensions):
    if dim.count(colon) == 1:
        print a[dim]

这个方法的结果是（我省略了一些简单的代码来打印左边的内容...）：

a[0,0,:] -->  [0 1 2 3]
a[0,1,:] -->  [4 5 6 7]
a[0,2,:] -->  [ 8  9 10 11]
a[0,:,0] -->  [0 4 8]
a[0,:,1] -->  [1 5 9]
a[0,:,2] -->  [ 2  6 10]
a[0,:,3] -->  [ 3  7 11]
a[1,0,:] -->  [12 13 14 15]
a[1,1,:] -->  [16 17 18 19]
a[1,2,:] -->  [20 21 22 23]
a[1,:,0] -->  [12 16 20]
a[1,:,1] -->  [13 17 21]
a[1,:,2] -->  [14 18 22]
a[1,:,3] -->  [15 19 23]
a[:,0,0] -->  [ 0 12]
a[:,0,1] -->  [ 1 13]
a[:,0,2] -->  [ 2 14]
a[:,0,3] -->  [ 3 15]
a[:,1,0] -->  [ 4 16]
a[:,1,1] -->  [ 5 17]
a[:,1,2] -->  [ 6 18]
a[:,1,3] -->  [ 7 19]
a[:,2,0] -->  [ 8 20]
a[:,2,1] -->  [ 9 21]
a[:,2,2] -->  [10 22]
a[:,2,3] -->  [11 23]

这里的关键是，用 a 来索引，比如 a[0,0,:]，其实和用 a[(0,0,slice(None))] 是一样的。 （这只是普通的 Python 切片，不是 numpy 特有的。如果你想验证这一点，可以写一个简单的类，只实现 __getitem__ 方法，然后打印出你在索引这个类的实例时传入的内容。）

所以，我们想要的是从 0 到 nx、0 到 ny、0 到 nz 等等的每一种可能组合，并且每个轴上都有一个 None。

不过，我们想要的是一维数组，所以需要过滤掉那些有多于或少于一个 None 的情况（也就是说，我们不想要 a[:,:,:]、a[0,:,:]、a[0,0,0] 等等）。

希望这样能让你明白一些...

编辑：我假设顺序并不重要... 如果你需要问题中列出的确切顺序，你就需要对这个做一些修改...

Answer 3

这里有一个简单的迭代器实现：

def iter1d(a):
    return itertools.chain.from_iterable(
        numpy.rollaxis(a, axis, a.ndim).reshape(-1, dim)
        for axis, dim in enumerate(a.shape))

这个迭代器会按照你在帖子中给出的顺序，输出子数组：

for x in iter1d(a):
    print x

打印结果是：

[ 0 12]
[ 1 13]
[ 2 14]
[ 3 15]
[ 4 16]
[ 5 17]
[ 6 18]
[ 7 19]
[ 8 20]
[ 9 21]
[10 22]
[11 23]
[0 4 8]
[1 5 9]
[ 2  6 10]
[ 3  7 11]
[12 16 20]
[13 17 21]
[14 18 22]
[15 19 23]
[0 1 2 3]
[4 5 6 7]
[ 8  9 10 11]
[12 13 14 15]
[16 17 18 19]
[20 21 22 23]

这里的关键是要遍历所有的维度，然后对每个维度进行调整，把数组变成一个二维数组，这样每一行就是你想要的一维子数组。