Python按filter分组求和proportion，寻求更优Pythonic实现方案

输入数据 #

presets = [{'proportion': 1, 'filter': {'tagger_mood': ['sad', 'party']}},
{'proportion': 1, 'filter': {'vocal_instrumental': 1}},
{'proportion': 1, 'filter': {'vocal_instrumental': 2}},
{'proportion': 1.1, 'filter': {'tagger_mood': ['sad', 'party']}},
{'proportion': 1.1, 'filter': {'vocal_instrumental': 1}},
{'proportion': 1.1, 'filter': {'vocal_instrumental': 2}}]

需求 #

按filter字段分组，对每组的proportion求和，最终得到如下结构的结果：

[
    {'proportion': 2.1, 'filter': {'tagger_mood': ['sad', 'party']}},
    {'proportion': 2.1, 'filter': {'vocal_instrumental': 1}},
    {'proportion': 2.1, 'filter': {'vocal_instrumental': 2}}
]

原实现代码 #

presets = [...]
merged_filter = []
merged_proportion = []
for preset in presets:
    if preset['filter'] not in merged_filter:
        merged_filter.append(preset['filter'])
        merged_proportion.append(preset['proportion'])
    else:
        merged_proportion[merged_filter.index(preset['filter'])] += preset['proportion']
print([{'proportion': p, 'filter': f} for p, f in zip(merged_proportion, merged_filter)])

优化方案 #

原实现的时间复杂度为O(n²)，因为每次判断preset['filter'] not in merged_filter和调用merged_filter.index()都需要遍历列表。以下是几种更高效、符合Python风格的实现方式：

方案1：使用collections.defaultdict（最优效率） #

由于字典无法直接作为字典的键，我们先将filter字典转换为可哈希的tuple（对键值对排序避免顺序影响），再用defaultdict累加求和：

from collections import defaultdict

sum_dict = defaultdict(float)

for item in presets:
    # 将filter转为可哈希的键
    filter_key = tuple(sorted(item['filter'].items()))
    sum_dict[filter_key] += item['proportion']

# 转换为目标格式
result = [{'proportion': val, 'filter': dict(key)} for key, val in sum_dict.items()]
print(result)

该方案时间复杂度为O(n)，效率远高于原实现，逻辑清晰易读。

方案2：使用itertools.groupby（简洁风格） #

groupby需要先按分组键排序，适合数据量不大的场景：

from itertools import groupby

# 定义分组键的转换函数
def get_filter_key(item):
    return tuple(sorted(item['filter'].items()))

# 先按分组键排序
presets_sorted = sorted(presets, key=get_filter_key)

# 分组求和
result = []
for key, group in groupby(presets_sorted, key=get_filter_key):
    total_proportion = sum(item['proportion'] for item in group)
    result.append({'proportion': total_proportion, 'filter': dict(key)})

print(result)

该方案代码简洁，但排序带来O(n log n)的时间复杂度，适合对性能要求不极致的场景。

方案3：使用普通字典（无额外依赖） #

如果不想导入模块，可直接用普通字典实现，逻辑与defaultdict一致：

sum_dict = {}

for item in presets:
    filter_key = tuple(sorted(item['filter'].items()))
    if filter_key in sum_dict:
        sum_dict[filter_key] += item['proportion']
    else:
        sum_dict[filter_key] = item['proportion']

result = [{'proportion': val, 'filter': dict(key)} for key, val in sum_dict.items()]
print(result)

该方案无需依赖任何模块，同样实现O(n)的时间复杂度，简洁高效。

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