Python: 如果前两个元素相同,则在嵌套列表中添加元素

Python: adding elements in nested list if the first two elements are the same

我有一个嵌套列表:

 a = [[1,2,3,4,5],[1,2,5,6,7],[1,2,5,2,1],[1,3,4,5,7],[3,4,1,2,3],[3,4,1,1,1]]

如果子列表的前两个元素相同,那么我想添加它们对应的第3、4、5个元素,同时保留前两个元素。对于上述情况,结果应为:

a = [[1,2,13,12,13],[1,3,4,5,7],[3,4,2,3,4]]

有办法吗?

谢谢。

您可以使用字典根据前两项对 sub-lists 进行分类,然后遍历这些项目并计算总和:

>>> d = {}
>>> for sub in a:
...    d.setdefault(tuple(sub[:2]),[]).append(sub[2:])
... 

>>> 
>>> [k+tuple(map(sum, zip(*v))) for k,v in d.items()]
[(1, 2, 13, 12, 13), (1, 3, 4, 5, 7), (3, 4, 2, 3, 4)]

我会创建一个字典,其中第一对夫妇作为键,值是列表或元组:

a=[[1,2,3,4,5],[1,2,5,6,7],[1,2,5,2,1],[1,3,4,5,7],[3,4,1,2,3],[3,4,1,1,1]]

#create empty dict
d = {}

# function for adding list b to lists a
def addlist(a,b): return [a[i]+b[i] for i in range(len(a))]  

# now iter through each item and add to existing or create a new record in dictionary
for i in a:
    d[tuple(i[:2])] = addlist(d.setdefault(tuple(i[:2]),[0,0,0]),i[2:])

#the wanted output would be
[list(k)+v for k,v in d.items()]

Pandas one-liner

pandasone-liner.

怎么样

导入和数据:

>>> import pandas as pd
>>> a = [[1,2,3,4,5],[1,2,5,6,7],[1,2,5,2,1],[1,3,4,5,7],[3,4,1,2,3],[3,4,1,1,1]]

Pandas 在行动:

>>> pd.DataFrame(a).groupby([0, 1]).sum().reset_index().values.tolist()
[[1, 2, 13, 12, 13], [1, 3, 4, 5, 7], [3, 4, 2, 3, 4]]

Step-by-step

制作数据框:

>>> df = pd.DataFrame(a)

   0  1  2  3  4
0  1  2  3  4  5
1  1  2  5  6  7
2  1  2  5  2  1
3  1  3  4  5  7
4  3  4  1  2  3
5  3  4  1  1  1

按前两列分组并对其他列求和:

>>> df2 = df.groupby([0, 1]).sum()
>>> df2

      2   3   4
0 1            
1 2  13  12  13
  3   4   5   7
3 4   2   3   4

groupby 方法按 01 列分组。结果是一个 GroupBy 对象。你需要 "do" 一些东西来形象化它。在这里,我们对组中的值求和。看看格式很好的 table:

可能会更清楚一些

这个索引是 MultiIndex:

MultiIndex(levels=[[1, 3], [2, 3, 4]],
           labels=[[0, 0, 1], [0, 1, 2]],
           names=[0, 1])

有两个级别,名称分别为 01。最左边的列是带有标签 13 的第一层,下一个是带有标签 23 和 [=26= 的第二层].第一层中的标签 1 用于前两行。所有其他标签仅用于一行。

展平 multi-index:

>> df3 = df2.reset_index()

   0  1   2   3   4
0  1  2  13  12  13
1  1  3   4   5   7
2  3  4   2   3   4

转换成列表:

>>> df3.values.tolist()
[[1, 2, 13, 12, 13], [1, 3, 4, 5, 7], [3, 4, 2, 3, 4]]

这是itertools.groupby相对straight-forward的应用。

下面是使用嵌套列表理解的方法。

from itertools import groupby
from operator import itemgetter

a = [[1,2,3,4,5],[1,2,5,6,7],[1,2,5,2,1],[1,3,4,5,7],[3,4,1,2,3],[3,4,1,1,1]]
expected = [[1,2,13,12,13],[1,3,4,5,7],[3,4,2,3,4]]
print(expected)

a = [list(k) + [sum(t) for t in zip(*[u[2:] for u in g])]
    for k, g in groupby(a, itemgetter(0, 1))]
print(a)

输出

[[1, 2, 13, 12, 13], [1, 3, 4, 5, 7], [3, 4, 2, 3, 4]]
[[1, 2, 13, 12, 13], [1, 3, 4, 5, 7], [3, 4, 2, 3, 4]]

请注意,这假定列表 a 已经具有以相同的 2 个相邻元素开头的 sub-lists。如果不是这种情况,那么应该对其进行排序,例如

a.sort(key=itemgetter(0, 1))

在运行上面的代码之前。

为了便于阅读和分析,将本质上相同的算法分开。

keyfunc = lambda seq: seq[:2]

a.sort(key=keyfunc)
new_a = []
for k, g in groupby(a, key=keyfunc):
    tails = [u[2:] for u in g]
    sums = [sum(t) for t in zip(*tails)]
    new_a.append(k + sums)

print(new_a)