總結numpy中陣列的不均勻片段

Question

給定一個包含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)) 

定時這兩個大i輸入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數組。一個與之相近的內置函數可能會比（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