python - 从生成器创建两个串联数组

Question

请考虑 Python 2.7 中的以下示例。我们有一个返回两个一维 numpy 数组的任意函数 f()。请注意，通常 f() 可能会返回不同大小的数组，并且大小可能取决于输入。

现在我们想在 f() 上调用 map 并将结果连接到两个单独的新数组中。

import numpy as np

def f(x):
    return np.arange(x),np.ones(x,dtype=int)   

inputs = np.arange(1,10)
result = map(f,inputs)
x = np.concatenate([i[0] for i in result]) 
y = np.concatenate([i[1] for i in result]) 

这给出了预期的结果。但是，由于结果可能会占用大量内存，因此最好通过调用 imap 而不是 map 来使用生成器。

from itertools import imap
result = imap(f,inputs)
x = np.concatenate([i[0] for i in result]) 
y = np.concatenate([i[1] for i in result]) 

但是，这会产生错误，因为在我们计算 y 时生成器是空的。

有没有办法只使用一次生成器并仍然创建这两个串联数组？我正在寻找没有 for 循环的解决方案，因为重复连接/附加数组效率很低。

提前致谢。

Answer 1

Is there a way to use the generator only once and still create these two concatenated arrays?

是的，可以用 tee 克隆生成器:

import itertools
a, b = itertools.tee(result)

x = np.concatenate([i[0] for i in a]) 
y = np.concatenate([i[1] for i in b]) 

---

但是，使用 tee 对您的情况下的内存使用没有帮助。上述解决方案需要 5 N 内存才能运行:

• N 用于在 tee 中缓存生成器，
• 2 N 用于 np.concatenate 调用中的列表理解，
• 2 N 用于串联数组。

显然，我们可以通过删除 tee 来做得更好:

x_acc = []
y_acc = []
for x_i, y_i in result:
    x_acc.append(x_i)
    y_acc.append(y_i)

x = np.concatenate(x_acc)
y = np.concatenate(y_acc)

这又减少了一个 N，剩下 4 N。更进一步意味着删除中间列表并预分配 x 和 y。请注意，您不需要知道数组的确切 大小，只需知道上限:

x = np.empty(capacity)
y = np.empty(capacity)
right = 0
for x_i, y_i in result:
    left = right
    right += len(x_i)  # == len(y_i)  
    x[left:right] = x_i
    y[left:right] = y_i

x = x[:right].copy()
y = y[:right].copy()

事实上，您甚至不需要上限。只需确保 x 和 y 足够大以容纳新项目:

for x_i, y_i in result:
    # ...
    if right >= len(x):
        # It would be slightly trickier for >1D, but the idea
        # remains the same: alter the 0-the dimension to fit 
        # the new item.
        new_capacity = max(right, len(x)) * 1.5
        x = x.resize(new_capacity)
        y = y.resize(new_capacity)