满足条件时列表中的值的总和

Posted 2021-03-27

我有一个python列表l包含类Element的实例：

class Element:
    def __init__(self, id, value):
        self.id = id
        self.value = value

l = [Element(1, 100), Element(1, 200), Element(2, 1), Element(3, 4), Element(3, 4)]

现在我想总结value类的所有Elements成员，如果他们的id等于获得此列表：

l = [Element(1, 300), Element(2, 1), Element(3, 8)]

什么是最pythonic的方式来做到这一点？

答案

（差不多？）itertools不能做什么。看看groupby：

from itertools import groupby
from operator import attrgetter


class Element:
    def __init__(self, id, value):
        self.id = id
        self.value = value
    def __repr__(self):  # kudos @mesejo
        return "Element({}, {})".format(self.id, self.value)

l = [Element(1, 100), Element(1, 200), Element(2, 1), Element(3, 4), Element(3, 4)]

l.sort(key=attrgetter('id'))  # if it is already sorted by 'id', comment-out

res = [Element(g, sum(sub.value for sub in k)) for g, k in groupby(l, key=attrgetter('id'))]

这导致：

print(res)   # [Element(1, 300), Element(2, 1), Element(3, 8)]

另一答案

一种方法是创建一个defaultdict，将id映射到值的总和。然后我们可以获取这些结果并使用它们来构建一个新的Elements列表。一种方法是使用starmap将该字典的项目映射到Element的参数

from collections import defaultdict
from itertools import starmap

class Element:
    def __init__(self, id, value):
        self.id = id
        self.value = value
    def __repr__(self):
        return "Element({}, {})".format(self.id, self.value)

l = [Element(1, 100), Element(1, 200), Element(2, 1), Element(3, 4), Element(3, 4)]

d = defaultdict(int)

for e in l:
    d[e.id] += e.value

print(list(starmap(Element, d.items())))
# [Element(1, 300), Element(2, 1), Element(3, 8)]

另一答案

您还可以使用set获得所需的结果，以获得唯一的ID和sum来计算总值。例如：

class Element:
    def __init__(self, id, value):
        self.id = id
        self.value = value

l = [Element(1, 100), Element(1, 200), Element(2, 1), Element(3, 4), Element(3, 4)]

ids = set(elem.id for elem in l)
totals = [Element(i, sum(elem.value for elem in l if elem.id == i)) for i in ids]
# [Element(1, 300), Element(2, 1), Element(3, 8)]