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。

显然，每个新的费用组都由0中的index_data表示;否则，不使用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. ])

矢量化循环使用Numpy

1 个答案: