arrays - 将3D numpy数组转换为3个索引的列表

因此，我有一个大型3D数据矩阵，例如10000X10000X1000，现在我要做的是遍历3D数据矩阵的每个元素，并将相同大小的2个不同矩阵的索引和值写入文件，例如一条线：

i j k val1 val2

我目前的工作是在3个嵌套循环中运行，并按以下方式打印它，例如2个小型3D数据矩阵的示例和方法：

import numpy as np


vv1= np.array([[[1,2,3],[2,3,4],[3,4,5]],
                [[4,5,6],[5,6,7],[6,7,8]],
                [[7,8,9],[8,9,10],[9,10,11]]])

vv2= np.array([[[1,2,3],[2,3,4],[3,4,5]],
                [[4,5,6],[5,6,7],[6,7,8]],
                [[7,8,9],[8,9,10],[9,10,11]]])

for x in range(vv1.shape[0]):
    for y in range(vv1.shape[1]):
        for z in range(vv1.shape[2]):
            print("{:} {:} {:} {:} {:}".format(x,y,z,vv1[x,y,z], vv2[x,y,z]))

这个简单的代码可以完成工作，但速度很慢。

我想到的另一种方法是创建一个一维长向量，每个条目将是3个索引值，然后对打印应用相同的逻辑，例如嵌套循环示例：

vv_ind = []

for x in range(vv1.shape[0]):
    for y in range(vv1.shape[1]):
        for z in range(vv1.shape[2]):
            vv_ind.append([x,y,z])

for elem in vv_ind:
    i = tuple(elem)
    print("{:} {:} {:} {:} {:}".format(*elem, vv1[i], vv2[i]))

给出所需的输出。

我的问题如下：


有没有更多的“ pythonic”方式来创建该索引列表？
关于最后的打印循环：

for elem in vv_ind:
    i = tuple(elem)
    print("{:} {:} {:} {:} {:}".format(*elem, vv1[i], vv2[i]))

有更有效的方法吗？


同样，这里给出的数组只是虚设的

希望能有所帮助

最佳答案

要创建索引列表，可以使用函数product：

from itertools import product

product(*3 * [range(3)]) # generator of indices

要么

product(range(3), range(3), range(3))

要么

from itertools import product, repeat

product(*repeat(range(3), 3))

您可以简化代码：

from itertools import product, repeat

for idx in product(*repeat(range(3), 3)):
    print(*idx, vv1[idx], vv2[idx])

正如评论中提到的@a_guest一样，我们可以使用np.ndindex(*vv1.shape)代替product(*repeat(range(3), 3))：