在numpy中对数组的不均匀段求和

Question

给定一个ndarray x和一个一维数组，其中包含维度为的连续切片的长度x，我想计算一个包含所有切片之和的新数组。例如，在两个维度中，第一个维度的总和为：

>>> lens = np.array([1, 3, 2])
array([1, 3, 2])
>>> x = np.arange(4 * lens.sum()).reshape((4, lens.sum())).astype(float)
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.]])
# I want to compute:
>>> result
array([[  0.,   6.,   9.],
       [  6.,  24.,  21.],
       [ 12.,  42.,  33.],
       [ 18.,  60.,  45.]])
# 0 = 0
# 6 = 1 + 2 + 3
# ...
# 45 = 22 + 23

我想到的两种方法是：

a）使用累积和花式索引：

def cumsum_method(x, lens):
    xc = x.cumsum(1)
    lc = lens.cumsum() - 1
    res = xc[:, lc]
    res[:, 1:] -= xc[:, lc[:-1]]
    return res

b）使用bincount并智能地生成适当的bin：

def bincount_method(x, lens):
    bins = np.arange(lens.size).repeat(lens) + \
        np.arange(x.shape[0])[:, None] * lens.size
    return np.bincount(bins.flat, weights=x.flat).reshape((-1, lens.size))

将这两个时间安排在较大的输入上会使cumsum方法的性能稍好一些：

>>> lens = np.random.randint(1, 100, 100)
>>> x = np.random.random((100000, lens.sum()))
>>> %timeit cumsum_method(x, lens)
1 loops, best of 3: 3 s per loop
>>> %timeit bincount_method(x, lens)
1 loops, best of 3: 3.9 s per loop

有没有一种我显然更有效的方法？似乎本机c调用会更快，因为它不需要分配cumsum或bins数组。numpy内置函数执行的操作可能会比（a）或（b）好。我无法通过搜索和浏览文档找到任何东西。

注意，这类似于此问题，但求和间隔不是规则的。

Answer 1

您可以使用np.add.reduceat：

>>> np.add.reduceat(x, [0, 1, 4], axis=1)
array([[  0.,   6.,   9.],
       [  6.,  24.,  21.],
       [ 12.,  42.,  33.],
       [ 18.,  60.,  45.]])

的索引列表[0, 1, 4]意思是：“总结切片0:1，1:4和4:”。你可以从生成这些值lens使用np.hstack(([0], lens[:-1])).cumsum()。

即使将的指标计算在内lens，一种reduceat方法也可能比替代方法快得多：

def reduceat_method(x, lens):
    i = np.hstack(([0], lens[:-1])).cumsum()
    return np.add.reduceat(x, i, axis=1)

lens = np.random.randint(1, 100, 100)
x = np.random.random((1000, lens.sum())

%timeit reduceat_method(x, lens)
# 100 loops, best of 3: 4.89 ms per loop

%timeit cumsum_method(x, lens)
# 10 loops, best of 3: 35.8 ms per loop

%timeit bincount_method(x, lens)
# 10 loops, best of 3: 43.6 ms per loop