Python 3 按foo_bar_bazzle结果分组统计数组元素的高性能优化方案

原实现的核心问题在于存储了所有分组的具体元素，对于长度达10^9的数组来说，内存占用会达到无法承受的量级，同时频繁的列表append操作也会带来额外性能开销。我们只需要分组的计数而非元素本身，因此可以从以下方向优化：

1. 直接统计分组计数，不存储元素 #

放弃存储每个分组的元素列表，改用字典直接记录每个分组的出现次数，内存复杂度从O(n)降至O(k)（k为不同分组的数量，远小于10^9），同时避免了列表append的开销。

基础实现（用普通字典） #

rslt_counts = {}
for i in arr:
    key = foo_bar_bazzle(i)
    # 若键不存在则初始化为0，再+1
    rslt_counts[key] = rslt_counts.get(key, 0) + 1

# 计算最终结果
return sum(count ** 2 for count in rslt_counts.values())

更简洁的实现（用defaultdict） #

借助collections.defaultdict可以简化计数逻辑：

from collections import defaultdict

rslt_counts = defaultdict(int)
for i in arr:
    rslt_counts[foo_bar_bazzle(i)] += 1

return sum(count ** 2 for count in rslt_counts.values())

2. 优化foo_bar_bazzle函数的执行效率 #

如果foo_bar_bazzle本身计算耗时较高，且数组中存在大量重复的x，可以用缓存复用函数结果：

from functools import lru_cache
from collections import defaultdict

# 给函数加缓存，相同输入直接返回已计算的结果
@lru_cache(maxsize=None)
def foo_bar_bazzle(x):
    # 原函数逻辑
    return result

rslt_counts = defaultdict(int)
for i in arr:
    rslt_counts[foo_bar_bazzle(i)] += 1

return sum(count ** 2 for count in rslt_counts.values())

注意：仅当foo_bar_bazzle是纯函数（相同输入始终返回相同输出）时，缓存才有效。

3. 超大数组的分块处理（可选） #

如果数组无法一次性加载到内存（比如从磁盘读取的超大文件），可以分批次读取并累计计数，核心逻辑不变，只是把遍历拆分为多个批次：

from collections import defaultdict

rslt_counts = defaultdict(int)

# 假设arr是按批次读取的生成器
for batch in get_arr_batches():
    for i in batch:
        rslt_counts[foo_bar_bazzle(i)] += 1

return sum(count ** 2 for count in rslt_counts.values())

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