使用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. ])