创建2D numpy数组的Pythonic方式

Question

我有一个函数 gen()，它返回一个包含 nElements 个浮点数的 numpy 数组。我想找一种更符合 Python 风格的（最好是一行代码）方法来做以下事情：

a = zeros((nSamples, nElements))
for i in xrange(nSamples):
     a[i,:] = gen()

这是一种实现方式：

a = array([gen() for i in xrange(nSamples)]).reshape((nSamples, nElements))

不过，这种方法速度稍慢，因为没有提前分配好 numpy 数组：

import time
from numpy import *

nSamples  = 100000
nElements = 100

start = time.time()
a = array([gen() for i in xrange(nSamples)]).reshape((nSamples, nElements))
print (time.time() - start)

start = time.time()
a = zeros((numSamples, nElements))
for i in xrange(numSamples):
    a[i,:] = gen()
print (time.time() - start)

输出：

1.82166719437
0.502261161804

那么，有没有办法在保持预先分配数组的情况下，实现同样的一行代码呢？

谢谢，
-Nate

Answer 1

这可能没有直接回答你的问题，但因为你在标题中提到了Pythonic，所以我想说一下。请理解，Pythonic并不一定意味着用“一行代码”或者最聪明、最简短的方式来做事情。恰恰相反，Pythonic的代码更注重清晰易懂。

就你的代码而言，我觉得：

a = zeros((nSamples, nElements))
for i in xrange(nSamples):
     a[i,:] = gen()

比下面这个：

a = array([gen() for i in xrange(nSamples)]).reshape((nSamples, nElements))

要清晰得多。因此，我不会说第二个更符合Pythonic的风格，可能反而不太符合。

Answer 2

我觉得这个可以满足你的需求：

a = vstack([ gen() for _ in xrange(nSamples) ])

因为我无法访问你的 gen 函数，所以无法进行时间测试。另外，这个方法（还有你的一行代码）在内存使用上没有你那个 for 循环版本那么友好。你的一行代码会把所有 gen() 的输出都存起来，然后再构建数组，而 for 循环只需要同时在内存中保留一个 gen() 的输出（还有那个 numpy 数组）。