リストによって指定された可変長のnumpy cumsum列

Question

how to cumsum nリストで指定された各列の連続する要素。 cumsumは次の行以降にリセットされます。例えば、レンズ= [2,3] ==>第2行のCUMSUM：[2：5]は、次の3行のCUMSUM [2]

import numpy as np 
lens = [3, 2] 
a = np.array(
     [[ 1, 2], 
     [ 1, 2], 
     [ 1, 2], 
     [ 1, 2], 
     [ 1, 2]]) 

が

np.array(
     [[ 1, 2], 
     [ 2, 4], 
     [ 3, 6], 
     [ 1, 2], 
     [ 2, 4]]) 

を付与しようとしループを避けるため

Answer 1

一つのオプションは、アレイ、CUMSUMを分割し、それらを組み合わせている：

np.concatenate(list(map(lambda a: np.cumsum(a, axis=0), np.array_split(a, np.cumsum(lens))))) 
#array([[1, 2], 
#  [2, 4], 
#  [3, 6], 
#  [1, 2], 
#  [2, 4]], dtype=int32) 

---

分割せずに別の選択肢と組み合わせは、以下のような特定のインデックスで和をリセットする補助配列を作成することである。

idx = np.cumsum([0] + lens)[:-1] 
aux = np.zeros_like(a) 
aux[idx[1:], :] = -np.add.reduceat(a, idx)[:-1] 
(a + aux).cumsum(0) 

#array([[1, 2], 
#  [2, 4], 
#  [3, 6], 
#  [1, 2], 
#  [2, 4]], dtype=int32) 

2つの方法がほぼ同じ速度である：

def split_concat(a): 
    return np.concatenate(list(map(lambda a: np.cumsum(a, axis=0), np.array_split(a, np.cumsum(lens))))) 

def reset_sum(a): 
    idx = np.cumsum([0] + lens)[:-1] 
    aux = np.zeros_like(a) 
    aux[idx[1:], :] = -np.add.reduceat(a, idx)[:-1] 
    return (a + aux).cumsum(0) 


lens = np.arange(1000) 
a = np.ones((lens.sum(), 2)) 
(reset_sum(a) == split_concat(a)).all() 
# True 

%timeit split_concat(a) 
# 12.8 ms ± 35.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) 

%timeit reset_sum(a) 
# 13.6 ms ± 87.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)