在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) 使用 cumsum 和花式索引：

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) 更好。通过搜索和查看文档，我找不到任何东西。

注意，这类似于this question，但求和间隔不规则。

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: 求和”。您可以使用np.hstack(([0], lens[:-1])).cumsum() 从lens 生成这些值。

即使考虑到来自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