python - 访问 3 维 numpy 数组中给定 x、y 位置的所有元素

Question

mat_a = np.random.random((5, 5))
mat_b = np.random.random((5, 5))
mat_c = np.random.random((5, 5))
bigmat = np.stack((mat_a, mat_b, mat_c)) # this is a 3, 5, 5 array

for (x, y, z), value in np.ndenumerate(bigmat):
    print (x, y, z)

在上面的例子中，我怎样才能循环，以便我只遍历 5 x 5 数组，并且在每个位置我得到 3 个值，即循环应该运行 25 次，每次，我得到一个包含 3 个值的数组(一个来自 mat_a、mat_b 和 mat_c 中的每一个)

• 编辑:请注意，我以后需要能够按位置访问元素，即如果 bigmat 被 reshape ，应该有一种方法可以根据特定的 y、z 访问元素

Answer 1

有一个函数可以生成给定形状的所有索引，ndindex。

for y,z in np.ndindex(bigmat.shape[1:]):
    print(y,z,bigmat[:,y,z])

0 0 [ 0 25 50]
0 1 [ 1 26 51]
0 2 [ 2 27 52]
0 3 [ 3 28 53]
0 4 [ 4 29 54]
1 0 [ 5 30 55]
1 1 [ 6 31 56]
...

对于像这样的简单情况，它并不比双 for range 循环容易多少。也不会更快；但你要求迭代。

另一个迭代器是 itertools.product(range(5),range(5))

从时间上看，产品还不错:

In [181]: timeit [bigmat[:,y,z] for y,z in itertools.product(range(5),range(5
     ...: ))]
10000 loops, best of 3: 26.5 µs per loop

In [191]: timeit [bigmat[:,y,z] for (y,z),v in np.ndenumerate(bigmat[0,...])]
     ...: 
10000 loops, best of 3: 61.9 µs per loop

转置和 reshape 是获取三元组列表(或数组)的最快方法 - 但它也没有给出索引:

In [198]: timeit list(bigmat.transpose(1,2,0).reshape(-1,3))
100000 loops, best of 3: 15.1 µs per loop

但同样的操作是从np.mgrid(或np.meshgrid)获取索引:

np.mgrid[0:5,0:5].transpose(1,2,0).reshape(-1,2)

(虽然这出奇的慢)