Question

给定一个ndarray x和一个包含x维度的连续切片长度的一维数组，我想计算一个包含所有切片之和的新数组。例如，在对维度1求和的二维中：

>>> 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或bin数组。一个接近这个的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:求和”。您可以使用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

numpy中数组的不均匀段

1 个答案: