使用Numpy向量化循環

Question

我有兩個1D Numpy數組index_data和charge，它們的長度相同並且分別包含整數和浮點數。我正在嘗試製作total_charge陣列，如下所示：

total_charge = [] 
for i in range(len(index_data)): 
    if index_data[i] == 0: 
     total_charge.append(charge[i]) 
    else: 
     total_charge[-1] += charge[i] 
total_charge = np.array(total_charge) 

我該如何對此進行矢量化？幫助我Numpy巫師，你是我唯一的希望。

Answer 1

index_data和charge的具體示例將有助於澄清預期輸入。但是，從閱讀您所擁有的內容，我認爲可以使用np.add.reduceat。

顯然，每個新的收費組都由index_data中的0表示;否則，不使用index_data中的值。 total_charge是各組收費的總和。

下面是如何使用reduceat的示例。

首先，一些演示數據：

In [105]: index_data 
Out[105]: array([0, 1, 1, 0, 1, 1, 1, 0, 1]) 

In [106]: charge 
Out[106]: array([ 1.5, 2. , 3. , 2.5, 1.5, 1. , 1. , 2. , 1. ]) 

zerolocs給其中0發生在index_data指數：

In [107]: zerolocs = where(index_data==0)[0] 

In [108]: zerolocs 
Out[108]: array([0, 3, 7]) 

使用np.add.reduceat總結充電組。

In [109]: total_charge = np.add.reduceat(charge, zerolocs) 

In [110]: total_charge 
Out[110]: array([ 6.5, 6. , 3. ])