如何基于元组内指定索引对元组列表进行分组？

问题描述 #

现有列表：

A = [('a',(1,2,3)),
     ('b',(2,4,5)),
     ('c',(2,3,2)),
     ('d',(5,3,2))]

需要基于每个元素内部元组的第2、第3个元素（即内部元组索引1和2的位置）进行分组，期望输出：

output = [[('a',(1,2,3))],
          [('b',(2,4,5))],
          [('c',(2,3,2)), ('d',(5,3,2))]]

现有代码问题 #

你当前的代码先提取内部元组生成列表B，导致丢失了外层元组的第一个元素（如'a'、'b'），后续需要重新映射，效率较低：

from operator import itemgetter
from itertools import groupby

B = [i[1] for i in A]

semi_output = [list(g) for _,g in 
                groupby(B,itemgetter(1,2))]

高效优化方案 #

直接对原列表A操作，无需额外生成中间列表。注意itertools.groupby仅对连续的相同键分组，因此需要先按分组键排序，再执行分组：

from operator import itemgetter
from itertools import groupby

# 按分组键排序，确保同组元素连续排列
sorted_A = sorted(A, key=lambda x: (x[1][1], x[1][2]))
# 直接以原列表元素分组，分组键为内部元组的第2、3个元素
output = [list(group) for _, group in groupby(sorted_A, key=lambda x: (x[1][1], x[1][2]))]

如果想用itemgetter简化分组键定义，可以写成：

group_key = lambda x: itemgetter(1, 2)(x[1])
sorted_A = sorted(A, key=group_key)
output = [list(group) for _, group in groupby(sorted_A, key=group_key)]

输出验证 #

运行上述代码后，output将完全符合期望的分组结果。

内容的提问来源于stack exchange，提问作者tcokyasar