高效地将3D Numpy数组重新塑形为1D列表加坐标向量

Question

我有一个大的嵌套数组a（256x256x256），我需要将其重组成一个元素如下的列表：

[ (i,j,k), a[i,j,k] ]

我目前的做法如下：

aflat = a.flatten().tolist()
coords = list(itertools.product(range(256), repeat=3))
thelist = [list(x) for x in zip(coords, aflat)]

这个代码可以运行，但速度比较慢。

通过不再运行时生成坐标向量并从文件中读取它们，我可能可以节省一秒钟左右的时间。然而，主要的减速似乎在最后一行，超过了6秒。

有没有更快的方法在Python中生成所需的数据结构？

- Marco Tompitak

---

看起来问题在于动态分配。您正在创建2^(8*3)个列表。您最初选择的数据类型似乎有问题。您肯定想要一些固定大小的东西，可以预先分配...寻找“struct pack”。 - P i

可能会有帮助：numpy.ndenumerate。 - wwii

Answer 1

正如@P-i所评论的，主要问题在于该代码创建了大量列表，Python花费了很多时间进行内存管理。为了消除这个问题，您可以使用numpy数组来预分配数据，并使用其repeat和tile函数生成i,j,k值:

# order='F' is important here so column-wise assignment can
# occur with a stride of 1.  Switching the order results
# in a significant performance hit.
coords = numpy.zeros([a.size,4],'d',order='F')

NI, NJ, NK = a.shape

# build columns for (i,j,k) tuples using repeat and tile
coords[:,0] = numpy.repeat(range(NI),NJ*NK)
coords[:,1] = numpy.tile(numpy.repeat(range(NJ),NK), NI)
coords[:,2] = numpy.tile(range(NK), NI*NJ)
coords[:,3] = a.flatten()

这将生成一个数组，其中每一行都是（i，j，k，value）。它假定您的原始数组按行主要排序（numpy中的C-ordered数组）。

在我的测试中，基于Python 3.5在2013年款MacBook Pro上进行了10次迭代，运行OP的转换需要约20秒，而使用此方法每次转换只需约8秒。

输出格式必须是列表，数组可以在最后一步转换为列表。但是，在我的测试中，这会将转换时间增加到每次转换13秒。

- sfstewman

---

有趣的想法，但这真的更快吗？我可能做错了什么，在ipython中尝试一些简单的基准测试，发现这种方法比我发布的要慢大约30%。 - Marco Tompitak

@MarcoTompitak，我的时间结果与你的一致。经过一番调查研究，我认为主要的性能问题在于内存管理，正如其中一位评论者所指出的那样。我更深入地思考了这个问题，并发现如果将输出保留在数组格式中，可以节省大量时间。 - sfstewman

太神奇了！我不知道我们计时差异的原因在哪里，但在我的电脑上，你的方法比我的快8-9倍，而不仅仅是2倍。我重新编写了代码的其余部分，尽可能地处理数组（在某些时候我需要将其转换为列表），这样运行速度更快。瓶颈已经从数据结构的生成转移到了我实际进行的分析，这对我来说足够好了。 - Marco Tompitak

Answer 2

为了进一步解释@wii上面的评论，您正在寻找np.ndenumerate。

通常，您会避免显式创建列表，而是使用迭代器。例如：

for (i,j,k), val in np.ndenumerate(your_3d_array):
    assert val == your_3d_array[i,j,k]

# Note that we also could have done:
for ind, val in np.ndenumerate(your_3d_array):
    assert val == your_3d_array[ind]

然而，如果你确实想要创建完整的中间列表，你可以使用：

list(np.ndenumerate(your_3d_array))

作为一个更完整的例子：

In [1]: import numpy as np

In [2]: x = np.arange(3*4*5).reshape(3, 4, 5)

In [3]: x
Out[7]:
array([[[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19]],

       [[20, 21, 22, 23, 24],
        [25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34],
        [35, 36, 37, 38, 39]],

       [[40, 41, 42, 43, 44],
        [45, 46, 47, 48, 49],
        [50, 51, 52, 53, 54],
        [55, 56, 57, 58, 59]]])

In  [4]: list(np.ndenumerate(x))
Out [4]: 
[((0, 0, 0), 0),
 ((0, 0, 1), 1),
 ((0, 0, 2), 2),
 ((0, 0, 3), 3),
...
 ((2, 3, 1), 56),
 ((2, 3, 2), 57),
 ((2, 3, 3), 58),
 ((2, 3, 4), 59)]

- Joe Kington

---

这绝对简单得多，但似乎并不更快。在一个256x256x256的数组上尝试，这比我的代码慢了20%以上。它还提供了一个元组列表，而不是一个列表的列表。正是在我的代码中将其转换为列表导致了很大的减速。如果我只留下zip()并使用（所以我得到一个元组列表），我已经获得了合理的性能（比ndenumerate()快2.5倍以上）。在你那边有什么不同吗？ - Marco Tompitak