python - 基于行输入的 Python 条件和

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]

---

如果第 3 方包可能适合您，您也可以使用 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}

或者直接使用groupedby的reduce参数:

>>> groupedby(zip(A, B), key=itemgetter(0), keep=lambda x: x[0]*x[1], reduce=add)
{1: 0.75, 2: 1.98, 3: 5.1}

免责声明:我是 iteration_utilities 的作者包。