Question

我试图用Python做一个条件求和产品。简化的想法如下：

A = [1 1 2 3 3 3]
B = [0.50 0.25 0.99 0.80 0.70 0.20]

我想输出

Total1 = 0.50*1 + 0.25*1
Total2 = 0.99*2
Total3 = 0.80*3 + 0.70*3 + 0.20*3

我正在考虑使用FOR ... IF ...结构，为A中的给定值指定B中的所有相应值应该相加。

实际上它是一个庞大的数据集，所以我必须让脚本能够遍历所有类别？

此刻，我正在努力将这个想法转化为适当的Python脚本。有人能指出我正确的方向吗？

Answer 1

这似乎非常适合itertools.groupby（假设A中的值已排序，它可能无法正常用于A=[1,1,2,2,1]）：

from itertools import groupby
A = [1, 1, 2, 3, 3, 3]
B = [0.50, 0.25, 0.99, 0.80, 0.70, 0.20]

for key, grp in groupby(zip(A, B), key=lambda x: x[0]):
    grp = [i[1] for i in grp]
    print(key, key * sum(grp))

打印：

1 0.75
2 1.98
3 5.1

您也可以将其存储在列表中，而不是打印值：

res = []
for key, grp in groupby(zip(A, B), key=lambda x: x[0]):
    grp = [i[1] for i in grp]
    res.append(key*sum(grp))
print(res)
# [0.75, 1.98, 5.1]

如果第三方套餐可能适合您，您也可以使用iteration_utilities.groupedby：

>>> from iteration_utilities import groupedby
>>> from operator import itemgetter, add

>>> {key: key*sum(value) for key, value in groupedby(zip(A, B), key=itemgetter(0), keep=itemgetter(1)).items()}
{1: 0.75, 2: 1.98, 3: 5.1}

或直接使用reduce的{{1}}参数：

groupedby

免责声明：我是iteration_utilities软件包的作者。

Answer 2

我认为您可以使用itertools.groupby解决此问题：

import itertools
from operator import itemgetter

results = [group * sum(v[1] for v in values)
           for group, values in itertools.groupby(zip(A, B), itemgetter(0))]

这假设A中所有相等的数字彼此相邻。如果它们可能不存在，您需要对它们进行排序或使用不同的算法。

Answer 3

如果您不介意使用numpy并假设订购了这些组，您可以通过以下方式进行：

A = [1, 1, 2, 3, 3, 3]
B = [0.50, 0.25, 0.99, 0.80, 0.70, 0.20]
A = np.asarray([1, 1, 2, 3, 3, 3])
B = np.asarray([0.50, 0.25, 0.99, 0.80, 0.70, 0.20])
index = np.full(len(A),True)
index[:-1] = A[1:] != A[:-1]
prods = A*B

#result
res = np.add.reduceat(prods, np.append([0], (np.where(index)[0]+1)[:-1]))

此外，如果您有大型列表，这可以真正加快操作

Python中基于行输入的条件求和

3 个答案: