如何使用itertools分组并保留指定键与名称信息？

需求说明 #

现有SQL查询返回的数据集：

[(2, 5, 'JOHN K', 'YAHOO'), (2, 6, 'AARON M', 'YAHOO'), (2, 7, 'NICK C', 'YAHOO'), (1, 2, 'CELESTE G', 'GOOGLE'), (1, 3, 'RICH M', 'GOOGLE'), (1, 4, 'SANDEEP C', 'GOOGLE')]

需要修改现有代码，实现：

• 分组键为**(公司名, 元组第0位数字)**（可接受类似'YAHOO,2'的字符串格式）
• 对应值为**(姓名, 元组第1位数字)**的拼接组合，格式如'(JOHN K,5);(AARON M,6);(NICK C,7)'

修改后的代码 #

import itertools
import operator


def accumulate(rows):
    # 按(公司名, 元组第0位)作为分组键
    it = itertools.groupby(rows, operator.itemgetter(3, 0))
    k = {}
    for key, subiter in it:
        # 将分组键转为字符串格式，如('YAHOO',2) -> 'YAHOO,2'
        str_key = f"{key[0]},{key[1]}"
        # 拼接每个项的(姓名, 元组第1位)为指定格式
        values = ';'.join(f'({item[2]},{item[1]})' for item in subiter)
        k[str_key] = values
    return k


if __name__ == '__main__':
    rows = [(2, 5, 'JOHN K', 'YAHOO'), (2, 6, 'AARON M', 'YAHOO'), (2, 7, 'NICK C', 'YAHOO'), 
            (1, 2, 'CELESTE G', 'GOOGLE'), (1, 3, 'RICH M', 'GOOGLE'), (1, 4, 'SANDEEP C', 'GOOGLE')] 
    groupedby = accumulate(rows)
    print(groupedby)

代码关键点说明 #

• 分组键修改：将operator.itemgetter(3)改为operator.itemgetter(3, 0)，同时按公司名（索引3）和元组第0位数字进行分组
• 键格式转换：把分组的元组键转为'公司名,数字'的字符串格式，符合期望输出要求
• 值拼接逻辑：遍历分组内的每个元素，将姓名（索引2）和元组第1位数字格式化为'(姓名,数字)'，再用分号拼接成字符串

如果希望分组键保留元组格式（如('YAHOO', 2)），只需去掉str_key的转换，直接用key作为字典的键即可：

# 替换键处理部分
k[key] = ';'.join(f'({item[2]},{item[1]})' for item in subiter)

此时输出会是：

{('YAHOO', 2): '(JOHN K,5);(AARON M,6);(NICK C,7)', ('GOOGLE', 1): '(CELESTE G,2);(RICH M,3);(SANDEEP C,4)'}

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