сумма неровных сегментов массива в 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 

два способа, которые приходят на ум:

а) Используйте cumsum и фантазии индексацию:

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

б) Используйте bincount и разумно сгенерировать соответствующие бункеры:

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 nput имел метод 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 builtin, которая делает что-то близкое к этому, может быть лучше, чем (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