对多个长度不均匀的 NumPy 向量求和的最快方法

Question

问题陈述很简单：给定任意数量的 NumPy 一维浮点向量，如下所示：

v1 = numpy.array([0, 0, 0.5, 0.5, 1, 1, 1, 1, 0, 0])
v2 = numpy.array([4, 4, 4, 5, 5, 0, 0])
v3 = numpy.array([1.1, 1.1, 1.2])
v4 = numpy.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 10])

求和它们的最快方法是什么？

many_vectors = [v1, v2, v3, v4]

使用直接求和函数将不起作用，因为它们的长度可以是任意不均匀的：

>>> result = sum(many_vectors)
ValueError: operands could not be broadcast together with shapes (10,) (7,)

相反，可以求助于pandas 库，该库将提供一个简单的fillna 参数来避免此问题。

 >>> pandas.DataFrame(v for v in many_vectors).fillna(0.0).sum().values
 array([ 5.1,  5.1,  5.7,  5.5,  6. ,  1. ,  1. ,  1. ,  0. ,  0. ,  0. ,
    0. ,  0. ,  0. ,  0. , 10. ])

但这可能不是最优化的处理方式，因为生产用例将包含大量数据。

In [9]: %timeit pandas.DataFrame(v for v in many_vectors).fillna(0.0).sum().values
1.16 ms ± 97.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Answer 1

方法#1

由于输入数组的大小和数组数量如此之大，我们需要提高内存效率，因此建议使用一个循环的数组，一次迭代地添加一个数组 -

many_vectors = [v1, v2, v3, v4] # list of all vectors

lens = [len(i) for i in many_vectors]
L = max(lens)
out = np.zeros(L)
for l,v in zip(lens,many_vectors):
    out[:l] += v

方法 #2

另一个几乎矢量化的 masking 从那些不规则形状的向量/数组的列表中生成一个规则的 2D 数组，然后沿列求和以获得最终输出 -

# Inspired by https://stackoverflow.com/a/38619350/ @Divakar
def stack1Darrs(v):
    lens = np.array([len(item) for item in v])
    mask = lens[:,None] > np.arange(lens.max())
    out_dtype = np.result_type(*[i.dtype for i in v])
    out = np.zeros(mask.shape,dtype=out_dtype)
    out[mask] = np.concatenate(v)
    return out

out = stack1Darrs(many_vectors).sum(0)

Answer 2

功劳归@Divakar。这个答案只会扩展和改进他的答案。特别是我重写了函数以匹配我们的样式指南并为它们计时。

有两种可能的方法：

方法#1

###############################################################################
def sum_vectors_with_padding_1(vectors):
    """Given an arbitrary amount of NumPy one-dimensional vectors of floats,
    do an element-wise sum, padding with 0 any that are shorter than the
    longest array (see https://stackoverflow.com/questions/56166217).
    """
    import numpy
    all_lengths = [len(i) for i in vectors]
    max_length  = max(all_lengths)
    out         = numpy.zeros(max_length)
    for l,v in zip(all_lengths, vectors): out[:l] += v
    return out

方法 #2

###############################################################################
def sum_vectors_with_padding_2(vectors):
    """Given an arbitrary amount of NumPy one-dimensional vectors of floats,
    do an element-wise sum, padding with 0 any that are shorter than the
    longest array (see https://stackoverflow.com/questions/56166217).
    """
    import numpy
    all_lengths = numpy.array([len(item) for item in vectors])
    mask        = all_lengths[:,None] > numpy.arange(all_lengths.max())
    out_dtype   = numpy.result_type(*[i.dtype for i in vectors])
    out         = numpy.zeros(mask.shape, dtype=out_dtype)
    out[mask]   = numpy.concatenate(vectors)
    return out.sum(axis=0)

时机

>>> v1 = numpy.array([0, 0, 0.5, 0.5, 1, 1, 1, 1, 0, 0])
>>> v2 = numpy.array([4, 4, 4, 5, 5, 0, 0])
>>> v3 = numpy.array([1.1, 1.1, 1.2])
>>> v4 = numpy.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 10])
>>> many_vectors = [v1, v2, v3, v4]
>>> %timeit sum_vectors_with_padding_1(many_vectors)
12 µs ± 645 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
>>> %timeit sum_vectors_with_padding_2(many_vectors)
22.6 µs ± 669 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

看来方法一更好！